Serialize user defined functions and table providers via protobuf

[timsaucer]

Is your feature request related to a problem or challenge? Please describe what you are trying to do.

PyLogicalPlan can currently only serialize or deserialize built in functions and table providers. It currently uses DefaultLogicalExtensionCodec in this function.

Users would like to be able to serialize and deserialize plans that include custom functions and table providers. One such example is the Iceberg python integration.

This topic was mentioned in the #datafusion-python channel in Discord.

Describe the solution you'd like

Create a logical extension codec that can process user defined functions and table providers. This likely means some work in the upstream repository.

One approach could consist of:

• Add encode and decode methods to the FFI_TableProvider, FFI_ScalarUDF and so on.
• Implement LogicalExtensionCodec on ForeignTableProvider, ForeignScalarUDF, and so on where the only methods implemented are those appropriate to table provide, udf, etc. That is for ForeignTableProvider it only supports try_encode_table_provider and try_decode_table_provider.
• Add a LogicalExtensionCodec to datafusion-python. For the table providers, iterate through all registered table providers and see if any of them successfully decode. The encode method is more straightforward.

Describe alternatives you've considered

The user could output the logical plan as a SQL command and then pass that along as a string.

Additional context

[colinmarc]

I explored the solution space a bit today, and I don't think this problem is really solvable with the APIs as they currently exist.

Just to be clear about what is actually needed, consider the following setup:

1. A python service, let's call it "frontend", builds a datafusion plan and wants to serialize it. The plan includes an iceberg scan, which is provided by IcebergTableProvider, exposed into python via the pycapsule interface (__datafusion_table_provider__).
2. A downstream rust service ("backend") wants to deserialize the plan.
3. A downstream python service ("logger") also wants to deserialize the plan.

LogicalExtensionCodec is not compatible with ForeignTableProvider

As stated above, the python bindings currently use DefaultLogicalExtensionCodec:

datafusion-python/src/sql/logical.rs

Line 187 in 2e1b713

let codec = DefaultLogicalExtensionCodec {};

However, this is only half of the problem. Imagine that I fork datafusion-python and change the code to use a custom LogicalExtensionCodec - thereby skipping the problem of how to expose the right API. My custom LogicalExtensionCodec would look something like this:

fn try_encode_table_provider(
    &self,
    table_ref: &TableReference,
    node: std::sync::Arc<dyn TableProvider>,
    buf: &mut Vec<u8>,
) -> datafusion::common::Result<()> {
    let source = node.as_any();

    // This downcast won't work on ForeignTableProvider.
    if let Some(iceberg_table) = source.downcast_ref::<IcebergTableProvider>() {
        todo!()
    }
}

The problem is that the downcast will fail, because the TableProvider is a ForeignTableProvider - the type information was lost during the "detour" through python.

Python type association is one-way

In step 3 above, we're in the "logger" python code, trying to deserialize the plan that the "frontend" python serialized.

Our decode implementation might look like:

fn try_decode_table_provider(
    &self,
    buf: &[u8],
    table_ref: &TableReference,
    schema: SchemaRef,
    ctx: &SessionContext,
) -> datafusion::common::Result<std::sync::Arc<dyn TableProvider>> {
    if let Ok(iceberg_data) = IcebergTableData::decode(buf) {
        let table_provider = IcebergTableProvider::new(...)
        // what do we do with the table provider now?
    }

This will work, but the python code has no way to turn the resulting provider back into the python wrapper type, since that information is lost. Therefore the plan can't really "roundtrip" in any meaningful sense.

encode can be stored on FFI_TableProvider, but not decode

As described above, we could somehow force an implementation of encode into FFI_TableProvider. (It would probably have to involve storing an Arc<dyn LogicalExtensionCodec> in private_data.) However, in order to iterate through and try possible "decoders", a global registry of LogicalExtensionCodecs would be needed somewhere.

[timsaucer]

I agree with your assessment.

I am starting to think your original suggestion was the correct one. I'm sorry I took a detour in the above approach.

I think what we want is a FFI_LogicalExtensionCodec and then in python we expose an interface to take n such objects. Then our python logical extension codec does a pass through each one, returning on the first success and trying the default codec in the end. Do you think that will meet your needs?

[timsaucer]

I've been spending some time today looking at this. I'm try to capture the exact use case and how it might work for the implementation.

Suppose we want to create a library IcebergDF. This library will be used in both Rust and Python implementations. Within this library we have IcebergTableProvider which is a struct that implements TableProvider. We also have IcebergLogicalCodec which implements LogicalExtensionCodec. For the codec, the only two methods it supports are try_encode_table_provider and try_decode_table_provider.

We have created a logical plan in datafusion-python that includes a table from an IcebergTableProvider instance. This was something along the lines of.

my_iceberg_table = IcebergTable(somevars)
ctx.register_table("my_table", my_iceberg_table)
df = ctx.table("my_table")
plan = df.logical_plan()

In this snippet we want to call plan.to_proto to get a serialized plan bytes that we can send along to some remote parts of our larger system.

What we want is the to_proto method to have an updated signature to be like

def to_proto(self, codecs: *LogicalExtensionCodec) -> bytes:

In this case we would pass plan.to_proto(my_codec) where my_codec is an instance of IcebergLogicalCodec.

IcebergLogicalCodec is a python object that has a function __datafusion_logical_extension_codec__ that will give the PyCapsule FFI_LogicalExtensionCodec. Now we run into a our first problem: Even though we are passing an Iceberg table provider created within the same library as this extension codec, it cannot tell. We have passed the FFI table provider object back and forth from the Iceberg library to df-python and back.

I have thought of two approaches, but open to others:

• My first thought was to add try_encode to FFI_TableProvider and also ForeignTableProvider. Then the IcebergLogicalCodec can do a downcast to ForeignTableProvider and call through the methods to try_encode. This would require then that IcebergTableProvider also implement a trait letting us know it can doe the encoding. This seems very round-about and easy to implement incorrectly.
• My second though, prompted by a LLM suggestion, was to add in an optional magic number to FFI_TableProvider. Then in IcebergLogicalCodec it can check to see if the magic number is the same, indicating that we have been passed back a pointer from the same providing library. If that magic number matches, then we know it is safe to access directly the inner Arc<dyn TableProvider> by IcebergLogicalCodec.

Supposing one of those two approaches works, then we should be a the point where we can encode our table provider for later processing.

Now for the second half, we want to decode this process in either another python work flow or in a rust work flow. Focusing on the python side, since I think the rust side would be even more straightforward, we want to make a corresponding update to our from_proto signature.

@staticmethod
def from_proto(ctx: SessionContext, data: bytes, codecs: *LogicalExtensionCodec) -> LogicalPlan:

In this case IcebergLogicalCodec already knows about IcebergTableProvider since it is specifically designed to encode/decode these table providers. In our datafusion-python implementation of from_proto on the rust side, we access via FFI the codec(s). Then we iterate through them until one is successful. If none are successful we try the default. Should they all fail, we return an error.

If this all makes sense to you, then it seems like the path forward, which I have bits and pieces of implemented in a branch on my computer are:

1. Add FFI_LogicalExtensionCodec which is probably fine to currently limit to table providers. I'd almost immediately want to extend it to udf, udaf, and udwf.
2. Update all of the FFI structs to include an optional magic number with a description of how to use these.
3. Update the to_proto and from_proto methods as describe above to iterate through provided codecs.
4. Create an example implementation in the datafusion-repo that demonstrates full end to end creation of a plan, storing it to bytes, decoding in a different session context, and executing the plan.

What do you think?

[colinmarc]

I have thought of two approaches...

Either of these seems plausible to me! It still seems like a shame that the LogicalExtensionCodec has to be written with ForeignTableProvider in mind, but I'm not sure I see a better way.

Perhaps we could "upstream" the magic number / ID thing to datafusion proper so that you only have to write LogicalExtensionCodec to consider the magic number, and not do the downcast_ref pattern at all. That would also be useful for the decode flow - it could be stored on the CustomTableScan proto representation and passed to try_decode_table_provider.

By the way, I have a PR up to change the encode/decode flow to deal with TableSource implementations, not TableProviders directly: apache/datafusion#16750. I don't think that should affect this work much.

[milenkovicm]

I believe this could unblock ballista support. I'm not 100% sure but it looks like step in right direction.