Repository navigation
Add example for writing an AnalyzerRule #10855
Description
Activity
Perhaps @goldmedal is interested in this / has some example code to share
I think we have an existing example showing how to create an AnalyzerRule in
However, it only works with the analyzer and the optimizer. Maybe we can enhance it after #10849 to show how to apply the custom rules for the end-to-end DataFusion query flow.
Currently, in my personal work, I just use it like
let ctx = SessionContext::new(); let new_state = ctx .state() .add_analyzer_rule(Arc::new(ModelAnalyzeRule::new())) .add_analyzer_rule(Arc::new(ModelGenerationRule::new())); let new_ctx = SessionContext::new_with_state(new_state); // create a plan to run a SQL query let df = new_ctx.sql("SELECT * FROM datafusion.default.orders").await?;
After the API updated, I guess it would be
let ctx = SessionContext::new(); ctx.add_analyzer_rule(Arc::new(ModelAnalyzeRule::new())) .add_analyzer_rule(Arc::new(ModelGenerationRule::new())); // create a plan to run a SQL query let df = ctx.sql("SELECT * FROM datafusion.default.orders").await?;
After #10849 is done, if no one else is working on it, I think I can help with it.
Reacted by Andrew Lamb- changed the title
[-]Add example for writing an `AnalyzerRule[/-][+]Add example for writing an `AnalyzerRule`[/+]on Jun 11, 2024 I have some time on a plane today that I may use to try and write up this example. I was inspired by some discussion I had this week
Hi @alamb, I just want to share how I use user-defined AnalyzerRule. As you know, I'm working on reimplementing the semantic layer engine, wren-engine, for LLM using DataFusion. The project is still a work in progress. However, I think I have finished the part of integration with DataFusion. I think it could be a nice use case for
AnalyzerRule.The basic concept of wren-engine is that user can define a virtual modeling layer to apply on their physical data.
For example, given a csv file registered in DataFusion context called,datafusion.public.customers. We can wrap a virtual table for it called,wrenai.default.customers_model.
User can query it likeselect * from wrenai.default.customers_model
Then , wren engine translate it to a physical query:
SELECT "customers_model"."city", "customers_model"."id", "customers_model"."state" FROM ( SELECT "customers_model"."city", "customers_model"."id", "customers_model"."state" FROM ( SELECT "datafusion"."public"."customers"."city" AS "city", "datafusion"."public"."customers"."id" AS "id", "datafusion"."public"."customers"."state" AS "state" FROM "public"."customers" ) AS "customers_model" ) AS "customers_model"
It's a simple use case to show how it works. We have other features, such as
relationshipbetween models or calculated fields in the model. I don't go into detail about them here. However, I implemented all of them throughAnalyzerRuleandUserDefinedLogicalNode.The related PR, Canner/wren-engine#613, is still under review. However, you can find the usage of
AnalyzerRulein rule.rs and the plan node in plan.rs. I also added some examples to show how to use the API or integrate with the DataFusion query flow.Thanks to DataFusion for providing such amazing features, allowing me to implement a more stable and structured converter for the semantic layer.
Reacted by Andrew LambThank you @goldmedal - this is a great hint. I have some basic analyzer rule coded up but maybe your example is a better idea. I have a PR in progress that I need to polish up and then will ping you for review
Reacted by Jax Liu
Is your feature request related to a problem or challenge?
We have an example for writing a user defined optimizer rule in
datafusion/datafusion/core/tests/user_defined/user_defined_plan.rs
Line 251 in 3773fb7
However, we don't have a corresponding example for writing a user defined AnalyzerRule which also means the APIs for using them are complicated (see #10849 for example)
Describe the solution you'd like
The idea example I think would be ti Add a file to https://github.com/apache/datafusion/tree/3773fb7fb54419f889e7d18b73e9eb48069eb08e/datafusion-examples
user_defined_analyzer.rsPerhaps the example could show how to create an
AnalyzerRulethat replaced an experssion likea / bwith a function callsafe_div(a, b)wheresave_divDescribe alternatives you've considered
No response
Additional context
No response