Syntax for Wasm-based UDAFs

CREATE [OR REPLACE] AGGREGATE function_name ( [parameter_list] )
  RETURNS { data_type [data_type_modifier] }
  WITH STATE <data_type | HANDLE>
  AS WASM FROM content_src
    [WITH WIT FROM content_src]
  INITIALIZE WITH udf_function_name
  ITERATE WITH udf_function_name
  MERGE WITH udf_function_name
  TERMINATE WITH udf_function_name
  [ SERIALIZE WITH udf_function_name
    DESERIALIZE WITH udf_function_name ]
  DEFINER = 'user'@

parameter_list:
  data_type [data_type_modifier [, ...] ] [, ...]

data_type_modifier:
  DEFAULT default_value | NOT NULL | NULL | COLLATE collation_name

contentSrc: BASE64 '"' <base64> '"                |
            HTTP '"' <url> '"' <cred> <config>    |
            LINK <linkName> <connStr>             |
            S3 <s3_config>                        |
            AZURE <azure_config>                  |
            GCS <gcs_config>

Wasm-based UDAFs

A Wasm-based UDAF can optionally pass its state by reference instead of value. For aggregation operations that require a state with a large amount of data, passing state by reference can improve performance. For example, a Wasm-based UDAF that maintains a hash table with thousands of entries. In this case, passing state by value necessitates the serialization of the hash table in and out of each Wasm function for every single row. Hence, passing state by reference substantially eliminates this overhead by allowing the Wasm function to internally maintain the data structure and pass a pointer (i.e. a HANDLE) between rows.

The life-cycle of the HANDLE state must be managed in the Wasm code. The implicit contract states that you need to create a unique handle for each call to the INITIALIZE function. At the same time, ITERATE calls can freely update the state associated with a handle or return a new handle. However, any state that is no longer accessible by the handle must be cleaned up, i.e., for any handle that is not returned, its state must be cleaned up by the functions. This applies to INITIALIZE, MERGE, TERMINATE, SERIALIZE, and DESERIALIZE functions

Here's an example. Two states are passed to the MERGE function. The user can reuse any of the provided states or create a new one. The following pseudo-code blocks demonstrate these use cases:

-- Reuse one of the existing states
def merge(a, b):
a = a + b
delete b
return a

-- Return a new state 
def merge(a, b):
c = a + b;
delete a
delete b
return c

Where a, b, and c are synonymous to states. In this example, any state that is no longer in use, is deleted (cleaned up).

Because each Wasm UDF runs in its own in-process sandbox, the UDAF needs to import all the Wasm functions (INITIALIZE, ITERATE, MERGE, TERMINATE, SERIALIZE, and DESERIALIZE) from the same Wasm module in a single CREATE AGGREGATE statement. Use the AS WASM clause in the function definition to import the Wasm function.

For information on how to import Wasm functions, refer to Create Wasm UDFs. Because Wasm functions are sandboxed, combining Wasm functions and PSQL functions in the same CREATE AGGREGATE statement is not supported.

Create UDAF

The following example creates a new UDAF named avg_udaf, which uses a RECORD state type and has the same behavior as the builtin AVG function. Note that before we can create the UDAF, each of the prerequisite user-defined scalar functions (UDFs) must be created, as the UDAF definition depends on their existence.

Create UDF Dependencies

The avg_udaf example UDAF depends on the following UDFs:

DELIMITER //
CREATE FUNCTION avg_init() RETURNS RECORD(s BIGINT, c BIGINT) AS
  BEGIN
    RETURN ROW(0, 0);
  END //
DELIMITER ;

DELIMITER //
CREATE FUNCTION avg_iter(state RECORD(s BIGINT, c BIGINT), value BIGINT) RETURNS RECORD(s BIGINT, c BIGINT) AS
  BEGIN
    RETURN ROW(state.s + value, state.c + 1);
  END //
DELIMITER ;

DELIMITER //
CREATE FUNCTION avg_merge(state1 RECORD(s BIGINT, c BIGINT), state2 RECORD(s BIGINT, c BIGINT)) RETURNS RECORD(s BIGINT, c BIGINT) AS
  BEGIN
    RETURN row(state1.s + state2.s, state1.c + state2.c);
  END //
DELIMITER ;

DELIMITER //
CREATE FUNCTION avg_terminate(state RECORD(s BIGINT, c BIGINT)) RETURNS BIGINT AS
  BEGIN
    RETURN state.s / state.c;
  END //
DELIMITER ;

Create UDAF Example

Once the UDF dependencies have been created, you can create the UDAF. Execute the following statement:

CREATE AGGREGATE avg_udaf(BIGINT) RETURNS BIGINT
  WITH STATE RECORD(s BIGINT, c BIGINT)
  INITIALIZE WITH avg_init
  ITERATE WITH avg_iter
  MERGE WITH avg_merge
  TERMINATE WITH avg_terminate;

After the UDAF has been successfully created, execute the following commands to try it:

CREATE TABLE t (i BIGINT);

INSERT INTO t VALUES (1), (2), (3), (4), (5);

SELECT avg_udaf(i) FROM t;

+-------------+
| avg_udaf(i) |
+-------------+
|           3 |
+-------------+
1 row in set

Create aggregate using Definer

CREATE AGGREGATE avg_udaf(BIGINT) RETURNS BIGINT
  WITH STATE RECORD(s BIGINT, c BIGINT)
  INITIALIZE WITH avg_init
  ITERATE WITH avg_iter
  MERGE WITH avg_merge
  TERMINATE WITH avg_terminate
  DEFINER = 'user'@'%';

To view a list of the functions and see the definer used:

SHOW FUNCTIONS;

+------------------------+--------------------------------+---------+-------------+---------------+-----+
| Functions_in_func_proc | Function Type                  | Definer | Data Format | Runtime Type  | Link|
+------------------------+--------------------------------+---------+-------------+---------------+-----+
| avg_init               | User Defined Function          | root@%  |             | PSQL          |     |
| avg_iter               | User Defined Function          | root@%  |             | PSQL          |     |
| avg_merge              | User Defined Function          | root@%  |             | PSQL          |     |
| avg_terminate          | User Defined Function          | root@%  |             | PSQL          |     | 
+------------------------+--------------------------------+---------+-------------+---------------+-----+

Create Wasm-based UDAF

The following example imports Wasm-based UDAFs from the Bloom Filters in SingleStore GitHub repository into SingleStore. Bloom filters are a space-efficient probabilistic data structure that can be used to test whether an element is a member of a set. This example performs the following operations:

• Import a Wasm-based UDAF and UDF.

• Generate a Bloom Filter from a column with string values.

• Run the Bloom Filter on another column.

This example imports the following modules from the repository:

• bloom_filter: A UDAF that generates a Bloom Filter from a column of string values.

• bloom_maybe_exists: A UDF that returns 0 if the argument string does not match the specified filter. If the argument string matches the filter, it may return 1.

• bloom_init_handle: Initializes a Bloom Filter with a predefined set of parameters.

• bloom_update_handle: Adds a value to the Bloom Filter.

• bloom_merge_handle: Merges two Bloom Filters, and returns the handle to the merged filter.

• bloom_serialize_handle: Serializes a Bloom Filter.

• bloom_deserialize_handle: Deserializes a Bloom Filter.

Refer to the lib.rs file in the repository for the source code.

This example uses Visual Studio Code with the Dev Containers extension.

1. Clone the Github repository:

   Shell

   Copy

   git clone https://github.com/singlestore-labs/singlestoredb-extension-bloom-filters.git

2. In the Visual Studio Code Command Palette, select Reopen in Container, and use the Dockerfile to build the container. You can also open the repository in a Rust container, and then manually install the WASI Cargo extension.

3. Run the following command in the terminal to verify that the Cargo CLI is installed:

   Shell

   Copy

   cargo -V

   cargo x.x.x

4. Run the following command in the terminal to compile the Wasm modules:

   Shell

   Copy

   cargo wasi build --release

   This command generates the extension.wasm file in the target/wasm32-wasi/release directory.

5. Import the UDAF and UDF modules into SingleStore.

   1. Encode the extension.wasm and extension.wit files to a base-64 encoded string using the following commands:

      Shell

      Copy

      base64 -w 0 extension.wasm
      base64 -w 0 extension.wit

      Copy the base-64 encoded string returned by each command from the standard output (stdout).

   2. Run the following SQL statement to import the bloom_filter() Wasm-based UDAF. Specify the base-64 encoded strings for the .wasm and .wit files in the AS WASM FROM BASE64 and WITH WIT FROM BASE64 clauses, respectively.

      SQL

      Copy

      CREATE OR REPLACE AGGREGATE bloom_filter(
          text CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci NOT NULL)
      RETURNS LONGBLOB NOT NULL
      WITH STATE HANDLE
      AS WASM FROM BASE64 '<base64-string-for-extension.wasm>'
      WITH WIT FROM BASE64 '<base64-string-for-extension.wit>'
      INITIALIZE WITH bloom_init_handle
      ITERATE WITH bloom_update_handle
      MERGE WITH bloom_merge_handle
      TERMINATE WITH bloom_serialize_handle
      SERIALIZE WITH bloom_serialize_handle
      DESERIALIZE WITH bloom_deserialize_handle;

   3. Run the following SQL statement to import the bloom_maybe_exists() Wasm-based UDF. Specify the base-64 encoded strings for the .wasm and .wit files in the AS WASM FROM BASE64 and WITH WIT FROM BASE64 clauses, respectively.

      SQL

      Copy

      CREATE OR REPLACE FUNCTION bloom_maybe_exists
      AS WASM FROM BASE64 '<base64-string-for-extension.wasm>'
      WITH WIT FROM BASE64 '<base64-string-for-extension.wit>';

6. Run the following SQL statements to create a data set:

   SQL

   Copy

   CREATE DATABASE dbTest;
   USE dbTest;
   CREATE TABLE tbSrc (tmp TEXT);
   INSERT INTO tbSrc VALUES ("Aron"), ("Jane"), ("Adam"), ("Millie"), ("Turner");
   CREATE TABLE tbTest (tmp TEXT);
   INSERT INTO tbTest VALUES ("Smith"), ("Millie"), ("Gray"), ("Jane"), ("Rose");

7. Run the following SQL statement to generate the Bloom Filter from the tbSrc.tmp column and store them in a variable named @bloomf:

   SQL

   Copy

   SELECT bloom_filter(tbSrc.tmp) FROM tbSrc INTO @bloomf;

8. Run the Bloom Filter created earlier on the tbTest.tmp column using the following SQL statement:

   SQL

   Copy

   SELECT tmp "Value", bloom_maybe_exists(@bloomf, tmp) "Exists in tbTest" FROM tbTest;

   +--------+------------------+
   | Value  | Exists in tbTest |
   +--------+------------------+
   | Rose   |                0 |
   | Millie |                1 |
   | Gray   |                0 |
   | Jane   |                1 |
   | Smith  |                0 |
   +--------+------------------+

Refer to Theta Sketch for more user-defined aggregate function examples.