Example for `SqlExecutor::logical_optimizer` to avoid unsupported arrow functions

[ahirner]

Sqlite doesn't support arrow_cast. Given a sqlite table with timestamp in milliseconds:

CREATE TABLE events (
  published INTEGER DEFAULT (CAST(unixepoch('subsec') * 1000 as INTEGER)) NOT NULL,
  sequence_id INTEGER PRIMARY KEY AUTOINCREMENT,
  message BLOB NOT NULL
) STRICT;

This query on the corresponding FederatedTableSource fails:

select arrow_cast(max(published), 'Timestamp(Millisecond, None)') from events limit 1;

The error is:

Error("Failed to convert query result to Arrow: Failed to extract row value: Invalid column type Integer at index: 0, name: arrow_cast(max(sqlite.events.published),Utf8("Timestamp(Millisecond, None)"))")

The sql in SQLExecutor::execute is:

SELECT CAST(max(`events`.`published`) AS TIMESTAMP) AS `arrow_cast(max(events.published),Utf8(\"Timestamp(Millisecond, None)\"))` FROM `events` LIMIT 1

Solution

#117 added functionality to correct too optimistic federation:

Extends SqlExecutor trait with logical_optimizer method, this can allow executor to hook into federation planning, allowing for rewriting LogicalPlan and even placement of FederationPlanNode. This is useful for avoiding federating nodes that are only part of datafusion eg. UDF, UDAF.. etc.

Maybe there are other ways to avoid the cast, but they seem less robust than rules specific to the sql dialect.

[ahirner]

The only plan that the optimizer of SqlExecutor::logical_optimizer() may receive is LogicalPlan::Extension(FederatedPlanNode):

datafusion-federation/datafusion-federation/src/sql/mod.rs

Lines 105 to 108 in e0496b3

	let mut plan = LogicalPlan::Extension(ext_node);
	if let Some(mut rewriter) = self.planner.executor.logical_optimizer() {
	plan = rewriter(plan)?;
	}

One can't mutate its plan because FederatedPlanNode::plan(&self) -> &LogicalPlan is not mut nor is .plan pub. One can't really recreate FederatedPlanNode because its .planner is non pub and SQLFederationPlanner is non pub.
One can't easily make a new rewrite because &dyn OptimizerConfig is not readily available in this stripped down callback.

I wonder how to apply this hook. Was it meant to call rewriter on the original plan of Optimizer::rewrite @trueleo ?

[trueleo]

@ahirner Thanks for bringing this up. Usage for this hook was planned because i ran into a similar issue with datafusion-json, but i never got around to writing it in my project. This was clubbed with some other refactors so I assumed FederatedPlanNode could be rewritten without whole lot of testing. This also means unfortunately I don't have answer for original issue yet. Maybe it's a good time to have this discussion on a separate thread

I will fix the issue by making it possible to recreate the nodes.

[ahirner]

Due to this issue, I think it's fair that noone uses this hook yet.

I will fix the issue by making it possible to recreate the nodes.

I see the opportunity to lift this logic into FederationProvider's Optimizer, re #117 (comment).

making it possible to recreate the nodes.

The signature could remaing the same, but with a plan not wrapped as extension.

[trueleo]

I see the opportunity to lift this logic into FederationProvider's Optimizer,

This part should be delegated to SQLExecutor in my opinion, so each executor can pick the sub plan. There are a lot of parallel to that trait and FederationProvider. It's almost like SQLExecutor: FederationProvider.