jonahgao commented on code in PR #11089: URL: https://github.com/apache/datafusion/pull/11089#discussion_r1666524349
########## datafusion-examples/examples/analyzer_rule.rs: ########## @@ -0,0 +1,200 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::array::{ArrayRef, Int32Array, RecordBatch, StringArray}; +use datafusion::prelude::SessionContext; +use datafusion_common::config::ConfigOptions; +use datafusion_common::tree_node::{Transformed, TreeNode}; +use datafusion_common::Result; +use datafusion_expr::{col, lit, Expr, LogicalPlan, LogicalPlanBuilder}; +use datafusion_optimizer::analyzer::AnalyzerRule; +use std::sync::{Arc, Mutex}; + +/// This example demonstrates how to add your own [`AnalyzerRule`] to +/// DataFusion. +/// +/// [`AnalyzerRule`]s transform [`LogicalPlan`]s prior to the DataFusion +/// optimization process, and can be used to change the plan's semantics (e.g. +/// output types). +/// +/// This example shows an `AnalyzerRule` which implements a simplistic of row +/// level access control scheme by introducing a filter to the query. +/// +/// See [optimizer_rule.rs] for an example of a optimizer rule +#[tokio::main] +pub async fn main() -> Result<()> { + // AnalyzerRules run before OptimizerRules. + // + // DataFusion includes several built in AnalyzerRules for tasks such as type + // coercion which change the types of expressions in the plan. Add our new + // rule to the context to run it during the analysis phase. + let rule = Arc::new(RowLevelAccessControl::new()); + let ctx = SessionContext::new(); + ctx.add_analyzer_rule(Arc::clone(&rule) as _); + + ctx.register_batch("employee", employee_batch())?; + + // Now, planning any SQL statement also invokes the AnalyzerRule + let plan = ctx + .sql("SELECT * FROM employee") + .await? + .into_optimized_plan()?; + + // Printing the query plan shows a filter has been added + // + // Filter: employee.position = Utf8("Engineer") + // TableScan: employee projection=[name, age, position] + println!("Logical Plan:\n\n{}\n", plan.display_indent()); + + // Execute the query, and indeed no Manager's are returned + // + // +-----------+-----+----------+ + // | name | age | position | + // +-----------+-----+----------+ + // | Andy | 11 | Engineer | + // | Oleks | 33 | Engineer | + // | Xiangpeng | 55 | Engineer | + // +-----------+-----+----------+ + ctx.sql("SELECT * FROM employee").await?.show().await?; + + // We can now change the access level to "Manager" and see the results + // + // +----------+-----+----------+ + // | name | age | position | + // +----------+-----+----------+ + // | Andrew | 22 | Manager | + // | Chunchun | 44 | Manager | + // +----------+-----+----------+ + rule.set_show_position("Manager"); + ctx.sql("SELECT * FROM employee").await?.show().await?; + + // The filters introduced by our AanalyzerRule are treated the same as any Review Comment: ```suggestion // The filters introduced by our AnalyzerRule are treated the same as any ``` ########## datafusion-examples/examples/analyzer_rule.rs: ########## @@ -0,0 +1,200 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::array::{ArrayRef, Int32Array, RecordBatch, StringArray}; +use datafusion::prelude::SessionContext; +use datafusion_common::config::ConfigOptions; +use datafusion_common::tree_node::{Transformed, TreeNode}; +use datafusion_common::Result; +use datafusion_expr::{col, lit, Expr, LogicalPlan, LogicalPlanBuilder}; +use datafusion_optimizer::analyzer::AnalyzerRule; +use std::sync::{Arc, Mutex}; + +/// This example demonstrates how to add your own [`AnalyzerRule`] to +/// DataFusion. +/// +/// [`AnalyzerRule`]s transform [`LogicalPlan`]s prior to the DataFusion +/// optimization process, and can be used to change the plan's semantics (e.g. +/// output types). +/// +/// This example shows an `AnalyzerRule` which implements a simplistic of row +/// level access control scheme by introducing a filter to the query. +/// +/// See [optimizer_rule.rs] for an example of a optimizer rule +#[tokio::main] +pub async fn main() -> Result<()> { + // AnalyzerRules run before OptimizerRules. + // + // DataFusion includes several built in AnalyzerRules for tasks such as type + // coercion which change the types of expressions in the plan. Add our new + // rule to the context to run it during the analysis phase. + let rule = Arc::new(RowLevelAccessControl::new()); + let ctx = SessionContext::new(); + ctx.add_analyzer_rule(Arc::clone(&rule) as _); + + ctx.register_batch("employee", employee_batch())?; + + // Now, planning any SQL statement also invokes the AnalyzerRule + let plan = ctx + .sql("SELECT * FROM employee") + .await? + .into_optimized_plan()?; + + // Printing the query plan shows a filter has been added + // + // Filter: employee.position = Utf8("Engineer") + // TableScan: employee projection=[name, age, position] + println!("Logical Plan:\n\n{}\n", plan.display_indent()); + + // Execute the query, and indeed no Manager's are returned + // + // +-----------+-----+----------+ + // | name | age | position | + // +-----------+-----+----------+ + // | Andy | 11 | Engineer | + // | Oleks | 33 | Engineer | + // | Xiangpeng | 55 | Engineer | + // +-----------+-----+----------+ + ctx.sql("SELECT * FROM employee").await?.show().await?; + + // We can now change the access level to "Manager" and see the results + // + // +----------+-----+----------+ + // | name | age | position | + // +----------+-----+----------+ + // | Andrew | 22 | Manager | + // | Chunchun | 44 | Manager | + // +----------+-----+----------+ + rule.set_show_position("Manager"); + ctx.sql("SELECT * FROM employee").await?.show().await?; + + // The filters introduced by our AanalyzerRule are treated the same as any + // other filter by the DataFusion optimizer, including predicate push down + // (including into scans), simplifications, and similar optimizations. + // + // For example adding another predicate to the query + let plan = ctx + .sql("SELECT * FROM employee WHERE age > 30") + .await? + .into_optimized_plan()?; + + // We can see the DataFusion Optimizer has combined the filters together + // when we print out the plan + // + // Filter: employee.age > Int32(30) AND employee.position = Utf8("Manager") + // TableScan: employee projection=[name, age, position] + println!("Logical Plan:\n\n{}\n", plan.display_indent()); + + Ok(()) +} + +/// Example AnalyzerRulw ule that implements a very basic "row level access Review Comment: ```suggestion /// Example AnalyzerRule that implements a very basic "row level access ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
