Is it possible to change the API for an operation class?

[kesmit13]

We are using the JSONGetItem operation for accessing JSON objects in SingleStoreDB, however, it only allows for integers and strings to be provided. Our JSON functions allow specifying a full path to an object in JSON as a series of ints and / or strings. It would be nice if the JSONGetItem would allow a tuple of ints / strings kind of like numpy does. You could then specify a path to a nested JSON object as:

tbl.json_col['key1', 'key2', 3, 'key3']

Currently, you have to do the following, which results in multiple calls to our JSON functions rather than a single one.

tbl.json_col['key1']['key2'][3]['key3']

I tried hacking in a tuple_of object into the index of the existing class, but there are apparently properties that get cached when the class in created, so it didn't work. If I also patch in the cached properties, it gets further, but I get the following error.

File ~/.pyenv/versions/3.11.3/lib/python3.11/site-packages/ibis/expr/rules.py:52, in <genexpr>(.0)
     50 @public
     51 def highest_precedence_shape(nodes):
---> 52     return max(node.output_shape for node in nodes)

AttributeError: 'tuple' object has no attribute 'output_shape'

I'm just wondering if there is a safe way to do this type of thing, without the possibility of breaking other backends using the existing class?

[kszucs]

Hi @kesmit13!

Indeed, constructing a single JSONGetItem(tbl.json_col, "key1.key2.3.key3") would be indeed preferable.

The rules system you tried to change is actually removed in favor of a new type annotation based system, this is already used in the main branch.

The simplest solution to achieve this would be to eagerly rewrite the operation on construction, the implementation is located at [JSONValue.getitem](https://github.com/ibis-project/ibis/blob/master/ibis/expr/types/json.py#L99). You can access the underlying operation object and depending on its type you can unwrap two JSONGetItem(JSONGetItem(...))operations into a singleJSONGetItem`, something like the following:

 def __getitem__(self, key):
    op = self.op()
    if isinstance(op, ops.JSONGetItem):
        node = ops.JSONGetItem(op.arg, key=f"{op.index}.{key}")
    else:
        node = ops.JSONGetItem(op, key)

    return node.to_expr()

This should provide the requested feature, although there might be backends not supporting nested JSON access. It would be nice to go over the backends with existing ibis implementation for JSONGetItem and check whether those backends would support it or not.
Even if they don't we can overcome the problem by splitting the key during expression compilation per backend.

[kesmit13]

I did find a solution that might be the "right" way to do this, although it does take a bit more work. I created a custom JSON data type that I return when mapping SingleStoreDB <=> Ibis data types. I then have to create SingleStoreDBJSONValue, SingleStoreDBJSONScalar, and SingleStoreDBJSONColumn objects that handle the __getitem__ implementation. The code for most of this is below. This might be the better way to do what I'm doing since I plan on adding other SingleStoreDB-specific method to JSON objects (and probably other types in the future as well). I just have to hope that the technique for looking up value classes doesn't change in Ibis.

class SingleStoreDBJSON(dt.JSON):

    scalar = 'SingleStoreDBJSONScalar'
    column = 'SingleStoreDBJSONColumn'


class SingleStoreDBJSONValue(ir.JSONValue):
    def __getitem__(self, key):
        return registry.SingleStoreDBJSONGetItem(self, key).to_expr()


class SingleStoreDBJSONScalar(ir.JSONScalar, SingleStoreDBJSONValue):
    pass


class SingleStoreDBJSONColumn(ir.JSONColumn, SingleStoreDBJSONValue):
    def __getitem__(self, key):
        return SingleStoreDBJSONValue.__getitem__(self, key)


# Patch in our JSON classes so they can be located
ir.SingleStoreDBJSONScalar = SingleStoreDBJSONScalar
ir.SingleStoreDBJSONColumn = SingleStoreDBJSONColumn

This still doesn't solve the tuple issue I'm seeing and I can't use the main branch version because I'm working on 6.1 compatibility, so I'm not sure how to get around that one just yet.

[kszucs]

Actually this should be compatible with both versions, since it doesn't interfere with the validation system.

Is SingleStoreDBJSON different from dt.JSON or just a workaround to construct the extended ir objects? If not then you can just update the dt.JSON.scalar and .column properties to point to the new value types.

I just have to hope that the technique for looking up value classes doesn't change in Ibis.

I'd say that it will be stable for while, at least cannot any reason to change it at the moment.

I plan on adding other SingleStoreDB-specific method to JSON objects

Could you please provide more details about them? Perhaps we would like to have them upstream as well, just haven't implemented yet.

I'm also planning to add easier extensibility to ibis and this can be motivating case. Could you give me an example how you would imagine/expect it to work?

[kesmit13]

I'm mostly just trying to extend existing objects. I was trying not to mess with the existing objects since it's possible someone could be using more than one backend and I didn't want to break the others. I'd like to make JSON objects act more like Python dicts / lists, so I was going to attempt to add more of the dunder methods and other API methods such as __contains__, __iter__ (scalars only), __mul__, __add__, __sub__, __div__, __len__, keys, etc.

Some of the above are also for our numeric vectors in JSON or binary columns. I realize that the mapping won't be perfect between JSON and Python since a JSON object can contain either a list or an object, and there is no way to separate them in a JSON object API.

[kesmit13]

My primary motivation is to get rid of some ugliness in our APIs. We have functions that handle all of those operations, but the function names get really long, and there are specific versions for each data type that I'd like to merge together. I'd just like to make them look and act more like the Python counterparts.

tbl.col.json_array_contains('...')    =>   '...' in tbl.col
tbl.col.json_keys()  =>  tbl.col.keys()
tbl.col.json_length()  =>  len(tbl.col)

tbl.col.vector_mul([a, b, c])   =>  tbl.col * [a, b, c]
tbl.col.vector_add([a, b, c])   =>   tbl.col + [a, b, c]
tbl.col.scalar_vector_mul(y)   =>   tbl.col * y
tbl.col.scalar_vector_add(y)   =>   tbl.col + y

[kesmit13]

I did discover a fatal flaw in my technique. Any function that has an output_dtype of dt.json has to be changed to return my customer JSON type. That means wrapping all JSON functions that are built into Ibis with my own versions. If there are only a few it's not a big deal, but it's something I have to be aware of.

[kesmit13]

After playing with this for a while, I think rewriting the __getitem__ method would be the easiest way. I'm going to add a new operation called JSONGetPath that takes a tuple of objects as the second argument. I don't want to put the values into a string key because that limits what can be used there (e.g., you could no longer use an expression there), but this will alleviate the need for a whole new type.