RisingWave UDF materialization for dbt

[james-johnston-thumbtack]

Today, dbt doesn't offer support out-of-the-box for user-defined functions (UDFs) for any database engine. Many dbt users have asked for this support over the years and come up with various solutions. One of the more recent solutions was proposed by Paul Brabban on his blog: Materialized UDFs in a dbt World. Paul provided a prototype custom dbt materialization for BigQuery UDFs that accompanied this blog post.

A big advantage of using a custom dbt materialization for UDFs is that the UDF becomes part of your dbt DAG. You can use all the normal commands for building and running the UDF model to create it. It will be created in the proper order with other models that depend on the UDF. dbt node selectors can be used to optionally choose it when running, since it's part of the DAG. Overall, this seems to be a better solution than dbt run hooks.

Of course, the prototype was for BigQuery, not RisingWave. So I thought it would be useful to make a version of this UDF custom materialization for RisingWave. Here is what it looks like to use:

File add_pair.sql creates a UDF named add_pair:

{{ config(
  materialized='udf',
  parameters=['a int, b int'],
  returns='int',
  language='sql',
  function_type='scalar'
) }}

a + b

If you then do dbt run -s add_pair, it will create the RisingWave UDF. You could then reference it using the ref keyword, similar to the examples on Paul's blog, linked above.

Now, interesting as this may be, I don't think this makes sense to make it an official part of RisingWave's dbt adapter, which is why I'm posting here and not as a pull request. The reason is that dbt is now planning to add out-of-the-box UDF support to dbt Core. It is part of the May 2025 roadmap post. Thus, in my opinion, it would be better for RisingWave to wait for official UDF support to be added to dbt, and then update the RisingWave dbt adapter to support it.

The preliminary specification shared by dbt is conceptually very similar to what Paul shared above: you'd put the body of the UDF in a dedicated .sql file in your project somewhere, and define the UDF parameters in an accompanying YAML file. Then, you'd reference it using ref from other models, and it becomes part of your DAG.

Anyway, until something official from dbt and RisingWave comes along, here's the custom materialization itself. The code is licensed under Apache 2.0, so feel free to use it.

Click to expand: macros/materializations/udf.sql

-- The RisingWave UDF materialization is heavily based off of the RisingWave view materialization in addition to
-- Paul Brabban's prototype:
-- https://github.com/risingwavelabs/dbt-risingwave/blob/b628cc06710d13396b77281f59429010dca8ced6/dbt/include/risingwave/macros/materializations/view.sql
-- https://github.com/brabster/dbt_materialized_udf/blob/3f9e1b895432436543d977cf62061a2d45f95c3b/macros/materializations/udf.sql
{% materialization udf, adapter="risingwave" %}

  -- Determine basic information about old and new database objects, like if anything already exists with the old name.
  {%- set full_refresh_mode = should_full_refresh() -%}
  -- check for any existing non-UDF relation
  {%- set old_relation = load_relation(this) -%}
  -- adapter.load_relation doesn't know about UDFs, so we need to check if the function exists manually
  {%- set function_exists_query -%}
    select count(*)
    from rw_catalog.rw_functions as rf
    inner join rw_catalog.rw_schemas as rs
      on rf.schema_id = rs.id
    where
      current_database() = '{{ this.database }}'
      and rs.name = '{{ this.schema }}'
      and rf.name = '{{ this.identifier }}'
      and rf.type != 'AGGREGATE'
  {%- endset -%}
  {%- set function_exists = run_query(function_exists_query).rows[0].values()[0] > 0 -%}
  {%- set aggregate_exists_query -%}
    select count(*)
    from rw_catalog.rw_functions as rf
    inner join rw_catalog.rw_schemas as rs
      on rf.schema_id = rs.id
    where
      current_database() = '{{ this.database }}'
      and rs.name = '{{ this.schema }}'
      and rf.name = '{{ this.identifier }}'
      and rf.type = 'AGGREGATE'
  {%- endset -%}
  {%- set aggregate_exists = run_query(aggregate_exists_query).rows[0].values()[0] > 0 -%}
  {%- set target_relation = this.incorporate(type='external') -%}

  -- Determine what the SQL would look like for creating the UDF.  We do this up front, in case there are any errors
  -- in the configuration of the UDF.

  {%- set parameters = config.require('parameters') -%}
  {%- set ret = config.require('returns') -%}
  {%- set language = config.require('language') -%}
  {%- set function_type = config.get('function_type', default='scalar') -%}
  {%- set remote_function = config.get('remote_function') -%}
  {%- set remote_url = config.get('remote_url') -%}
  {%- set options = config.get('options', default={}) -%}
  -- Only some parameter values are supported.
  -- NOTE: there is also 'wasm' language, but embedding WebAssembly binaries seems like a handful to figure out, so
  -- we're skipping it for now.
  {% if language not in ['javascript', 'python', 'remote', 'rust', 'sql'] %}
    {{ exceptions.raise_compiler_error("Unsupported UDF language: " ~ language) }}
  {% endif %}
  {% if function_type not in ['aggregate', 'scalar', 'table'] %}
    {{ exceptions.raise_compiler_error("Unsupported UDF function_type: " ~ function_type) }}
  {% endif %}
  {% if language != 'remote' and (remote_url or remote_function) %}
    {{ exceptions.raise_compiler_error("remote_url or remote_function are only supported for remote language") }}
  {% endif %}
  {% if language == 'remote' and (not remote_url or not remote_function) %}
    {{ exceptions.raise_compiler_error("remote_url and remote_function are required for remote language") }}
  {% endif %}

  -- Render the CREATE FUNCTION statement.
  {%- set create_sql -%}
    create {{ 'function' if function_type != 'aggregate' else 'aggregate' }} {{ target_relation }}(
      {{ parameters | join(', ') }}
    )
    {% if function_type == 'table' %}
      returns table ({{ ret | join(', ') }})
    {% else %}
      returns {{ ret }}
    {% endif %}
    {% if language != 'remote' %}
      language {{ language }}
    {% endif %}
    {% if remote_url %}
      as {{ adapter.quote(remote_function) }}
      using link '{{ remote_url | replace("'", "''") }}'
    {% else %}
      as '{{ sql | replace("'", "''") }}'
    {% endif %}
    {% if options %}
      with (
        {% for option_name, option_value in options.items() %}
          {{ option_name }} = {% if option_value is string %}
            '{{ option_value | replace("'", "''") }}'
          {% else %}
            {{ option_value }}
          {% endif %}
          {% if not loop.last %}, {% endif %}
        {% endfor %}
      )
    {% endif %}
    ;
  {%- endset -%}

  -- Drop the old relation or old function only if full refresh is requested.  This is consistent with the behavior of
  -- the other RisingWave adapter materializations.
  {% if full_refresh_mode %}
    {% if old_relation %}
      {{ adapter.drop_relation(old_relation) }}
    {% elif function_exists %}
      {%- set drop_sql -%}
        DROP FUNCTION {{ target_relation }} CASCADE;
      {%- endset -%}
      {% do run_query(drop_sql) %}
    {% elif aggregate_exists %}
      {%- set drop_sql -%}
        DROP AGGREGATE {{ target_relation }} CASCADE;
      {%- endset -%}
      {% do run_query(drop_sql) %}
    {% endif %}
    {%- set old_relation = none %}
    {%- set function_exists = false %}
    {%- set aggregate_exists = false %}
  {% endif %}

  -- Now, actually create the UDF.

  {{ run_hooks(pre_hooks, inside_transaction=False) }}
  {{ run_hooks(pre_hooks, inside_transaction=True) }}

  {% if old_relation is none and not function_exists and not aggregate_exists %}
    {% call statement('main') -%}
      {{ create_sql }}
    {%- endcall %}
  {% else %}
    {{ risingwave__execute_no_op(target_relation) }}
  {% endif %}

  -- NOTE: RisingWave doesn't yet support COMMENT ON statement for functions, so we skip the persist_docs call.

  {{ run_hooks(post_hooks, inside_transaction=False) }}
  {{ run_hooks(post_hooks, inside_transaction=True) }}

  {{ return({'relations': []}) }}

{% endmaterialization %}

And here is the documentation for how to use it:

Click to expand: macros/materializations/udf.yml

version: 2

macros:
  - name: materialization_udf_risingwave
    description: |
      Custom materialization providing support for UDF as a materialization type.

      This materialization is copied from and inspired by Paul Brabban's UDF materialization prototype:
      - [Source code](https://github.com/brabster/dbt_materialized_udf/blob/3f9e1b895432436543d977cf62061a2d45f95c3b/macros/materializations/udf.sql)
      - [Main blog post](https://tempered.works/posts/2024/02/19/materialized-udfs-in-a-dbt-world/)
      - [Another blog post](https://www.equalexperts.com/blog/tech-focus/udf-nirvana-a-new-way-to-integrate-udfs-into-the-dbt-ecosystem/)

      dbt Core's roadmap includes plans for first-class support for UDFs.  The
      [proposed spec](https://github.com/dbt-labs/dbt-core/discussions/11851) is very similar to Paul's prototype.
      Therefore, we expect that the UDF-as-a-materialization support found here will likely be very easy to migrate to
      future dbt Core support for UDFs.

      The key benefit of the materialization is that the UDFs become a first-class node in the dbt DAG.  This allows us
      to properly canary and deploy changes to UDFs as part of the regular RisingWave deployment procedures, for
      example.

      The RisingWave `udf` materialization has the following configs that can be provided via the standard dbt `config`
      block or Jinja macro:

      | Config name | Type | Required | Description |
      |-------------|------|----------|-------------|
      | `parameters` | `list[str]` | **Yes** | List of input parameters for the UDF.  Each parameter is a string in the format `name data_type`. |
      | `returns` | `str` | **Yes** | Return data type of the UDF. |
      | `language` | `str` | **Yes** | Language of the UDF.  Supported values for RisingWave are `javascript`, `python`, `remote`, `rust`, and `sql`. |
      | `function_type` | `str` | **No** | Type of the UDF.  Supported values are `scalar`, `aggregate`, and `table`.  `scalar` functions accept scalar inputs and return a scalar value.  `aggregate` functions aggregate multiple scalar inputs and return a single scalar output.  `table` functions accept scalar inputs and return a table.  Default is `scalar`. |
      | `remote_function` | `str` | **No** | Name of the remote function to call on the UDF server.  Required if `language` is `remote`. |
      | `remote_url` | `str` | **No** | URL of the UDF server.  Required if `language` is `remote`. |
      | `options` | `dict[str, any]` | **No** | Dictionary of UDF options to pass in the `WITH` clause.  Values are converted from Python data types to appropriate SQL data types. |

      The body of the dbt data model is used as the contents of the UDF.  If the UDF has `language` of `remote`, then it
      is a remote UDF, and the body is ignored.

      NOTE: the `RETURNS` clause for SQL UDFs is unsupported.  Users can prefix the expression they would have used
      with `SELECT`, and achieve the same thing.  In fact, that's exactly what the RisingWave parser does when it
      encounters a `RETURNS` clause.