This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 59e2c7e024 Minor: Add examples of using TreeNode with Expr (#10686)
59e2c7e024 is described below
commit 59e2c7e0246f219f6a9786b3f77873c51744e4ff
Author: Andrew Lamb <[email protected]>
AuthorDate: Wed May 29 06:39:24 2024 -0400
Minor: Add examples of using TreeNode with Expr (#10686)
---
datafusion/expr/src/expr.rs | 57 +++++++++++++++++++++++++++++++++++++++++++++
1 file changed, 57 insertions(+)
diff --git a/datafusion/expr/src/expr.rs b/datafusion/expr/src/expr.rs
index 3542e2d985..71cf3adddf 100644
--- a/datafusion/expr/src/expr.rs
+++ b/datafusion/expr/src/expr.rs
@@ -62,6 +62,13 @@ use sqlparser::ast::NullTreatment;
/// See [`ExprSchemable::get_type`] to access the [`DataType`] and nullability
/// of an `Expr`.
///
+/// # Visiting and Rewriting `Expr`s
+///
+/// The `Expr` struct implements the [`TreeNode`] trait for walking and
+/// rewriting expressions. For example [`TreeNode::apply`] recursively visits
an
+/// `Expr` and [`TreeNode::transform`] can be used to rewrite an expression.
See
+/// the examples below and [`TreeNode`] for more information.
+///
/// # Examples
///
/// ## Column references and literals
@@ -156,6 +163,56 @@ use sqlparser::ast::NullTreatment;
/// Expr::from(Column::from_qualified_name("t1.c2")),
/// ]);
/// ```
+///
+/// # Visiting and Rewriting `Expr`s
+///
+/// Here is an example that finds all literals in an `Expr` tree:
+/// ```
+/// # use std::collections::{HashSet};
+/// use datafusion_common::ScalarValue;
+/// # use datafusion_expr::{col, Expr, lit};
+/// use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
+/// // Expression a = 5 AND b = 6
+/// let expr = col("a").eq(lit(5)) & col("b").eq(lit(6));
+/// // find all literals in a HashMap
+/// let mut scalars = HashSet::new();
+/// // apply recursively visits all nodes in the expression tree
+/// expr.apply(|e| {
+/// if let Expr::Literal(scalar) = e {
+/// scalars.insert(scalar);
+/// }
+/// // The return value controls whether to continue visiting the tree
+/// Ok(TreeNodeRecursion::Continue)
+/// }).unwrap();;
+/// // All subtrees have been visited and literals found
+/// assert_eq!(scalars.len(), 2);
+/// assert!(scalars.contains(&ScalarValue::Int32(Some(5))));
+/// assert!(scalars.contains(&ScalarValue::Int32(Some(6))));
+/// ```
+///
+/// Rewrite an expression, replacing references to column "a" in an
+/// to the literal `42`:
+///
+/// ```
+/// # use datafusion_common::tree_node::{Transformed, TreeNode};
+/// # use datafusion_expr::{col, Expr, lit};
+/// // expression a = 5 AND b = 6
+/// let expr = col("a").eq(lit(5)).and(col("b").eq(lit(6)));
+/// // rewrite all references to column "a" to the literal 42
+/// let rewritten = expr.transform(|e| {
+/// if let Expr::Column(c) = &e {
+/// if &c.name == "a" {
+/// // return Transformed::yes to indicate the node was changed
+/// return Ok(Transformed::yes(lit(42)))
+/// }
+/// }
+/// // return Transformed::no to indicate the node was not changed
+/// Ok(Transformed::no(e))
+/// }).unwrap();
+/// // The expression has been rewritten
+/// assert!(rewritten.transformed);
+/// // to 42 = 5 AND b = 6
+/// assert_eq!(rewritten.data, lit(42).eq(lit(5)).and(col("b").eq(lit(6))));
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub enum Expr {
/// An expression with a specific name.
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
