How to extend Backends with non-scalar functions

[saschahofmann]

In previous versions, the docs described a way to add functions to a Backend. For example, here a reminder on how it was done in Ibis 4.

Luckily, the approach still works and I only had to update some of the class attribute's names and definitions after upgrading to Ibis 8 🎉

import ibis.expr.datatypes as dt
import ibis.expr.rules as rlz
from ibis.backends.bigquery.compiler import BigQueryExprTranslator

compiles = BigQueryExprTranslator.compiles

class JSONExtract(Value):
    value: Value[dt.String]
    json_path: Value[dt.String]
    shape = rlz.shape_like("value")
    dtype = dt.string


def json_extract(value, json_path):
    return JSONExtract(value, json_path).to_expr()


StringValue.json_extract = json_extract


@compiles(JSONExtract)
def _json_extract(t, expr):
    value, json_path = expr.op().args
    t_value = t.translate(value)
    t_json_path = t.translate(json_path)

    return f"JSON_QUERY({t_value}, {t_json_path})"

Now I would like to add the same functionality for postgres but the PostgreSQLExprTranslator doesn't have a compiles function? Is there a new, shared way to extend backends? The docs mention the scalar.builtin decorator but what about functions that don't return scalars?

I'd also be happy to hear whether I could use sqlglot to maybe add a backend agnostic solution (or a solution that works where ever its available).

[cpcloud]

@saschahofmann Thanks for raising this.

We removed that section from the docs due to the instability of the internal operations APIs.

Still, I'm glad to hear that it wasn't too bad to update your custom operations.

In your particular case, we do support that operation on JSON typed columns, which you should be able to access from string columns by casting:

In [1]: from ibis.interactive import *

In [2]: t = ibis.literal('{"a":1}')

In [3]: t
Out[3]: '{"a":1}'

In [4]: t.cast("json")
Out[4]: {'a': 1}

In [5]: t.cast("json")["a"]
Out[5]: 1

In [6]: type(t.cast("json")["a"])
Out[6]: ibis.expr.types.json.JSONScalar

In [7]: ibis.to_sql(t.cast("json")["a"])
Out[7]:
SELECT
  CAST('{"a":1}' AS JSON) -> '$.a' AS "JSONGetItem(Cast('{""a"":1}', json), 'a')"

In [8]: ibis.to_sql(t.cast("json")["a"], dialect="postgres")
Out[8]:
SELECT
  JSON_EXTRACT_PATH(CAST('{"a":1}' AS JSON), 'a') AS "JSONGetItem(Cast('{""a"":1}', json), 'a')"

Regarding scalar.builtin, the name scalar here is mean't to distinguish between a function that runs conceptually on every row of its inputs and produces one row of output per call. It does not mean that the return type of calling the function on an ibis expression will be an Ibis scalar expression. The naming comes from the convention used in some of our various backends.

We're considering ways to support custom operations in the newly refactored sqlglot-ized version of Ibis, so stay tuned!

[saschahofmann]

I see. So scalar is more in contrast to e.g. a reduction?

Right now, I have a bunch of functions that I added for BigQuery. A lot of them are around substracting times or dates from each other using e.g. BigQuery's TIMESTAMP_DIFF that excepts two columns and a unit to specify which unit(e.g. year, month or day difference) to extract. I can see that in the meantime ibis implemented a lot of functionality that would probably allow me to use a combination of functions to achieve this backend agnostic.

Just out of interest, imagine a theoretical case where both BigQuery and Postgres support a certain functionality but maybe have slightly different syntax e.g. a different function name (maybe in Postgres TIMESTAMP_DIFF is called TIMESTAMP_DIFFERENCE). Is there a way to use scalar.builtin to specify to use one of them depending on the backend.

Something along these lines ?

@ibis.udf.scalar.builtin(name="timestamp_diff", backend="bigquery")
@ibis.udf.scalar.builtin(name="timestamp_difference", backend="postgres")
def timestamp_diff(a: dt.datetime, b: dt.datetime) -> float:
   ...

What would be the way to extend Postgres similar to how I did it for BigQuery above? I remember it used to be the case that BigQuery was a "string" backend and Postgres "SQLAlchemy" but now all backends are "sqlglot" backends right?

In any case, I will see whether I can remove all my custom definitions 🙏

[cpcloud]

#8748 is a related issue.

I think what you're asking for--ability to define your own Value nodes and corresponding compiler rules--is something we're interested in tackling. Not 100% sure on what timeline, but probably in the next quarter or two if I had to guess.

[cpcloud]

For BigQuery's timestamp_diff and corresponding date version you should be able to use the ibis delta method.

[saschahofmann]

Alright, let's see, 🤞 I wont need the extra compilers anymore. Thanks again.

We hope to share the project with the public soon as well 😊

[saschahofmann]

I got another question: This time I am trying to find a generic way to implement bigquery's ANY_VALUE that returns a single row from a reduction. Postgres 16 seems to have a similar functionality.

As before, for BigQuery I am using this to compile bigquery code

class AnyValue(Reduction):
    arg: Value

    dtype = rlz.dtype_like("arg")


def any_value(arg):
    return AnyValue(arg).to_expr()


Column.any_value = any_value


@compiles(AnyValue)
def _any_value(t, expr):
    (arg,) = expr.op().args

    return f"ANY_VALUE({t.translate(arg)})"

I am using this to create a distinct functionality across several columns. You group by the columns you want to be distinct and select any value from the others. Another solution would be to have the distinct method accept columns but I don't think all backends support this.