Enable custom aggregate functions (take 2) (#529)

* initial commit.

* documentation

* remove no-longer-valid type

* close over state initialization for performance

* link documentation in comment

* more testing

* run tests if they're main

* accept a single arg

* this kind of works but I'm abandoning this branch

Basically it seems that the sqlite extension pattern of 'allocate
a struct and stick it in the context pointer' is not going to work
for us here. I wonder if using the id of the pointer returned by
sqlite3_aggregate_context would be enough? Since no two functions
could use the same pointer, per https://www.sqlite.org/c3ref/aggregate_context.html ?

* a middle road sqlite3_agg_context solution

* try out auto-updating state

* improve quantile test, add multiple agg test

* add a null to the test

* acorn fails to parse ||=, whatever

* make eslint happy

* make initial_value an argument

* test step and finalize exceptions

* add memory leak test

* update docs to current interface

* delete state in exception handlers

* remove null state

* return init function and document object

* more tests and update back to init function

* update redefinition test for new interface

* update README to match fixed signature

* more consistent test formatting

* Update README.md

Co-authored-by: Ophir LOJKINE <[email protected]>

* clarify what exactly the result will contain

* Update README.md

* Update README.md

* Update README.md

* Update README.md

* Update README.md

* Improve documentation and type annotations

* ignore documentation in eslintrc

* reduce code size

Thanks a lot, @llimllib !
 
Co-authored-by: dogquery <>
Co-authored-by: Ophir LOJKINE <[email protected]>

Author: llimllib

.eslintrc.js

@@ -16,6 +16,7 @@ module.exports = {
     ignorePatterns: [
         "/dist/",
         "/examples/",
+        "/documentation/",
         "/node_modules/",
         "/out/",
         "/src/shell-post.js",

README.md

@@ -75,6 +75,33 @@ db.create_function("add_js", add);
 // Run a query in which the function is used
 db.run("INSERT INTO hello VALUES (add_js(7, 3), add_js('Hello ', 'world'));"); // Inserts 10 and 'Hello world'
 
+// You can create custom aggregation functions, by passing a name
+// and a set of functions to `db.create_aggregate`:
+//
+// - an `init` function. This function receives no argument and returns
+//   the initial value for the state of the aggregate function.
+// - a `step` function. This function takes two arguments
+//    - the current state of the aggregation
+//    - a new value to aggregate to the state
+//  It should return a new value for the state.
+// - a `finalize` function. This function receives a state object, and
+//   returns the final value of the aggregate. It can be omitted, in which case
+//   the final value of the state will be returned directly by the aggregate function.
+//
+// Here is an example aggregation function, `json_agg`, which will collect all
+// input values and return them as a JSON array:
+db.create_aggregate(
+  "json_agg",
+  {
+    init: () => [],
+    step: (state, val) => [...state, val],
+    finalize: (state) => JSON.stringify(state),
+  }
+);
+
+db.exec("SELECT json_agg(column1) FROM (VALUES ('hello'), ('world'))");
+// -> The result of the query is the string '["hello","world"]'
+
 // Export the database to an Uint8Array containing the SQLite database file
 const binaryArray = db.export();
 ```

src/api.js

@@ -225,6 +225,14 @@ Module["onRuntimeInitialized"] = function onRuntimeInitialized() {
         "",
         ["number", "string", "number"]
     );
+
+    // https://www.sqlite.org/c3ref/aggregate_context.html
+    // void *sqlite3_aggregate_context(sqlite3_context*, int nBytes)
+    var sqlite3_aggregate_context = cwrap(
+        "sqlite3_aggregate_context",
+        "number",
+        ["number", "number"]
+    );
     var registerExtensionFunctions = cwrap(
         "RegisterExtensionFunctions",
         "number",
@@ -1131,81 +1139,90 @@ Module["onRuntimeInitialized"] = function onRuntimeInitialized() {
         return sqlite3_changes(this.db);
     };
 
-    /** Register a custom function with SQLite
-    @example Register a simple function
-        db.create_function("addOne", function (x) {return x+1;})
-        db.exec("SELECT addOne(1)") // = 2
+    var extract_blob = function extract_blob(ptr) {
+        var size = sqlite3_value_bytes(ptr);
+        var blob_ptr = sqlite3_value_blob(ptr);
+        var blob_arg = new Uint8Array(size);
+        for (var j = 0; j < size; j += 1) {
+            blob_arg[j] = HEAP8[blob_ptr + j];
+        }
+        return blob_arg;
+    };
 
-    @param {string} name the name of the function as referenced in
-    SQL statements.
-    @param {function} func the actual function to be executed.
-    @return {Database} The database object. Useful for method chaining
-     */
+    var parseFunctionArguments = function parseFunctionArguments(argc, argv) {
+        var args = [];
+        for (var i = 0; i < argc; i += 1) {
+            var value_ptr = getValue(argv + (4 * i), "i32");
+            var value_type = sqlite3_value_type(value_ptr);
+            var arg;
+            if (
+                value_type === SQLITE_INTEGER
+                || value_type === SQLITE_FLOAT
+            ) {
+                arg = sqlite3_value_double(value_ptr);
+            } else if (value_type === SQLITE_TEXT) {
+                arg = sqlite3_value_text(value_ptr);
+            } else if (value_type === SQLITE_BLOB) {
+                arg = extract_blob(value_ptr);
+            } else arg = null;
+            args.push(arg);
+        }
+        return args;
+    };
+    var setFunctionResult = function setFunctionResult(cx, result) {
+        switch (typeof result) {
+            case "boolean":
+                sqlite3_result_int(cx, result ? 1 : 0);
+                break;
+            case "number":
+                sqlite3_result_double(cx, result);
+                break;
+            case "string":
+                sqlite3_result_text(cx, result, -1, -1);
+                break;
+            case "object":
+                if (result === null) {
+                    sqlite3_result_null(cx);
+                } else if (result.length != null) {
+                    var blobptr = allocate(result, ALLOC_NORMAL);
+                    sqlite3_result_blob(cx, blobptr, result.length, -1);
+                    _free(blobptr);
+                } else {
+                    sqlite3_result_error(cx, (
+                        "Wrong API use : tried to return a value "
+                        + "of an unknown type (" + result + ")."
+                    ), -1);
+                }
+                break;
+            default:
+                sqlite3_result_null(cx);
+        }
+    };
+
+    /** Register a custom function with SQLite
+      @example <caption>Register a simple function</caption>
+          db.create_function("addOne", function (x) {return x+1;})
+          db.exec("SELECT addOne(1)") // = 2
+
+      @param {string} name the name of the function as referenced in
+      SQL statements.
+      @param {function} func the actual function to be executed.
+      @return {Database} The database object. Useful for method chaining
+       */
     Database.prototype["create_function"] = function create_function(
         name,
         func
     ) {
         function wrapped_func(cx, argc, argv) {
+            var args = parseFunctionArguments(argc, argv);
             var result;
-            function extract_blob(ptr) {
-                var size = sqlite3_value_bytes(ptr);
-                var blob_ptr = sqlite3_value_blob(ptr);
-                var blob_arg = new Uint8Array(size);
-                for (var j = 0; j < size; j += 1) {
-                    blob_arg[j] = HEAP8[blob_ptr + j];
-                }
-                return blob_arg;
-            }
-            var args = [];
-            for (var i = 0; i < argc; i += 1) {
-                var value_ptr = getValue(argv + (4 * i), "i32");
-                var value_type = sqlite3_value_type(value_ptr);
-                var arg;
-                if (
-                    value_type === SQLITE_INTEGER
-                    || value_type === SQLITE_FLOAT
-                ) {
-                    arg = sqlite3_value_double(value_ptr);
-                } else if (value_type === SQLITE_TEXT) {
-                    arg = sqlite3_value_text(value_ptr);
-                } else if (value_type === SQLITE_BLOB) {
-                    arg = extract_blob(value_ptr);
-                } else arg = null;
-                args.push(arg);
-            }
             try {
                 result = func.apply(null, args);
             } catch (error) {
                 sqlite3_result_error(cx, error, -1);
                 return;
             }
-            switch (typeof result) {
-                case "boolean":
-                    sqlite3_result_int(cx, result ? 1 : 0);
-                    break;
-                case "number":
-                    sqlite3_result_double(cx, result);
-                    break;
-                case "string":
-                    sqlite3_result_text(cx, result, -1, -1);
-                    break;
-                case "object":
-                    if (result === null) {
-                        sqlite3_result_null(cx);
-                    } else if (result.length != null) {
-                        var blobptr = allocate(result, ALLOC_NORMAL);
-                        sqlite3_result_blob(cx, blobptr, result.length, -1);
-                        _free(blobptr);
-                    } else {
-                        sqlite3_result_error(cx, (
-                            "Wrong API use : tried to return a value "
-                            + "of an unknown type (" + result + ")."
-                        ), -1);
-                    }
-                    break;
-                default:
-                    sqlite3_result_null(cx);
-            }
+            setFunctionResult(cx, result);
         }
         if (Object.prototype.hasOwnProperty.call(this.functions, name)) {
             removeFunction(this.functions[name]);
@@ -1229,6 +1246,137 @@ Module["onRuntimeInitialized"] = function onRuntimeInitialized() {
         return this;
     };
 
+    /** Register a custom aggregate with SQLite
+      @example <caption>Register a custom sum function</caption>
+        db.create_aggregate("js_sum", {
+            init: () => 0,
+            step: (state, value) => state + value,
+            finalize: state => state
+        });
+        db.exec("SELECT js_sum(column1) FROM (VALUES (1), (2))"); // = 3
+
+      @param {string} name the name of the aggregate as referenced in
+      SQL statements.
+      @param {object} aggregateFunctions
+                      object containing at least a step function.
+      @param {function(): T} [aggregateFunctions.init = ()=>null]
+            a function receiving no arguments and returning an initial
+            value for the aggregate function. The initial value will be
+            null if this key is omitted.
+      @param {function(T, any) : T} aggregateFunctions.step
+            a function receiving the current state and a value to aggregate
+            and returning a new state.
+            Will receive the value from init for the first step.
+      @param {function(T): any} [aggregateFunctions.finalize = (state)=>state]
+            a function returning the result of the aggregate function
+            given its final state.
+            If omitted, the value returned by the last step
+            will be used as the final value.
+      @return {Database} The database object. Useful for method chaining
+      @template T
+       */
+    Database.prototype["create_aggregate"] = function create_aggregate(
+        name,
+        aggregateFunctions
+    ) {
+        // Default initializer and finalizer
+        var init = aggregateFunctions["init"]
+            || function init() { return null; };
+        var finalize = aggregateFunctions["finalize"]
+            || function finalize(state) { return state; };
+        var step = aggregateFunctions["step"];
+
+        if (!step) {
+            throw "An aggregate function must have a step function in " + name;
+        }
+
+        // state is a state object; we'll use the pointer p to serve as the
+        // key for where we hold our state so that multiple invocations of
+        // this function never step on each other
+        var state = {};
+
+        function wrapped_step(cx, argc, argv) {
+            // > The first time the sqlite3_aggregate_context(C,N) routine is
+            // > called for a particular aggregate function, SQLite allocates N
+            // > bytes of memory, zeroes out that memory, and returns a pointer
+            // > to the new memory.
+            //
+            // We're going to use that pointer as a key to our state array,
+            // since using sqlite3_aggregate_context as it's meant to be used
+            // through webassembly seems to be very difficult. Just allocate
+            // one byte.
+            var p = sqlite3_aggregate_context(cx, 1);
+
+            // If this is the first invocation of wrapped_step, call `init`
+            //
+            // Make sure that every path through the step and finalize
+            // functions deletes the value state[p] when it's done so we don't
+            // leak memory and possibly stomp the init value of future calls
+            if (!Object.hasOwnProperty.call(state, p)) state[p] = init();
+
+            var args = parseFunctionArguments(argc, argv);
+            var mergedArgs = [state[p]].concat(args);
+            try {
+                state[p] = step.apply(null, mergedArgs);
+            } catch (error) {
+                delete state[p];
+                sqlite3_result_error(cx, error, -1);
+            }
+        }
+
+        function wrapped_finalize(cx) {
+            var result;
+            var p = sqlite3_aggregate_context(cx, 1);
+            try {
+                result = finalize(state[p]);
+            } catch (error) {
+                delete state[p];
+                sqlite3_result_error(cx, error, -1);
+                return;
+            }
+            setFunctionResult(cx, result);
+            delete state[p];
+        }
+
+        if (Object.hasOwnProperty.call(this.functions, name)) {
+            removeFunction(this.functions[name]);
+            delete this.functions[name];
+        }
+        var finalize_name = name + "__finalize";
+        if (Object.hasOwnProperty.call(this.functions, finalize_name)) {
+            removeFunction(this.functions[finalize_name]);
+            delete this.functions[finalize_name];
+        }
+        // The signature of the wrapped function is :
+        // void wrapped(sqlite3_context *db, int argc, sqlite3_value **argv)
+        var step_ptr = addFunction(wrapped_step, "viii");
+
+        // The signature of the wrapped function is :
+        // void wrapped(sqlite3_context *db)
+        var finalize_ptr = addFunction(wrapped_finalize, "vi");
+        this.functions[name] = step_ptr;
+        this.functions[finalize_name] = finalize_ptr;
+
+        // passing null to the sixth parameter defines this as an aggregate
+        // function
+        //
+        // > An aggregate SQL function requires an implementation of xStep and
+        // > xFinal and NULL pointer must be passed for xFunc.
+        // - http://www.sqlite.org/c3ref/create_function.html
+        this.handleError(sqlite3_create_function_v2(
+            this.db,
+            name,
+            step.length - 1,
+            SQLITE_UTF8,
+            0,
+            0,
+            step_ptr,
+            finalize_ptr,
+            0
+        ));
+        return this;
+    };
+
     // export Database to Module
     Module.Database = Database;
 };

src/exported_functions.json

@@ -41,5 +41,6 @@
 "_sqlite3_result_int",
 "_sqlite3_result_int64",
 "_sqlite3_result_error",
+"_sqlite3_aggregate_context",
 "_RegisterExtensionFunctions"
 ]