[Question] Converting TableSource to custom TableProvider

[nevi-me]

When registering a custom TableProvider, the provider gets converted to a dyn TableSource during logical planning. When executing the plan, the source is cast back to a provider using source_to_provider.

/// Attempt to downcast a TableSource to DefaultTableSource and access the
/// TableProvider. This will only work with a TableSource created by DataFusion.
pub fn source_as_provider(
    source: &Arc<dyn TableSource>,
) -> datafusion_common::Result<Arc<dyn TableProvider>> {
    match source
        .as_ref()
        .as_any()
        .downcast_ref::<DefaultTableSource>()
    {
        Some(source) => Ok(source.table_provider.clone()),
        _ => Err(DataFusionError::Internal(
            "TableSource was not DefaultTableSource".to_string(),
        )),
    }
}

When executing the physical plan, the custom TableProvider was previously used to scan() and return data. So the above function doesn't work as it only works with a DefaultTableSource.

It's unclear how I can avoid this by replacing with a custom call that'd return my custom table provider. How can I achieve this?

For context, I have the repo https://github.com/nevi-me/datafusion-rdbms-ext which has a PostgresTableProvider. It used to work before the recent changes to add a TableSource, but now I'm stuck and I can't figure out an upgrade path.

CC @andygrove re the TableSource change.

I'm looking in some repos that use DataFusion, but haven't found something yet.

[andygrove]

Hi @nevi-me where do you create the logical plan in this repo? Can you create a DefaultTableSource wrapping your provider?

[nevi-me]

Hi @nevi-me where do you create the logical plan in this repo? Can you create a DefaultTableSource wrapping your provider?

If I understand your question correctly, I don't construct any LogicalPlan with the provider/source.
I register all DB tables in the catalog, and the TableProvider is used by the default planner to create a LogicalPlan::TableScan with my table provider (now table source).

Then when I'm optimizing the plan, I would rewrite the table scan to an extension (https://github.com/nevi-me/datafusion-rdbms-ext/blob/main/src/optimizer.rs#L65)

Before the TableSource change, I would downcast dyn TableProvider to my PostgresTableProvider in plan_extension() https://github.com/nevi-me/datafusion-rdbms-ext/blob/main/src/physical_plan.rs#L209-L219.

So now I get the table source back, but because it's created with the default provider_as_source, I can't retrieve my custom table provider, which has SQL connection details.

[nevi-me]

I have a commit which makes provider_as_source part of TableProvider (7b73e0d)

trait TableProvider {
   fn as_source(self: Arc<Self>) -> Arc<dyn TableSource>;
}

This would solve one half of my problem, as I could create a SqlTableSource that also has the connection details that I need.

The part that's not straightforward is the corresponding TableSource::as_provider(), as it creates a circular dependency between the crates.

[nevi-me]

Can you create a DefaultTableSource wrapping your provider?

@andygrove this works, so I'll use that. I downcast to DefautTableSource, then downcast to PostgresTableProvider

[andygrove]

The part that's not straightforward is the corresponding TableSource::as_provider(), as it creates a circular dependency between the crates.

Yes, the reason for introducing TableSource in the first place was to stop logical plans depending on execution plans (via the TableProvider::scan method) and to allow us to move logical plans and expressions out of the core crate. It is intentionally difficult to reintroduce that dependency.

[andygrove]

I think we can close this as complete now.