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?