Skip to content

Commit

Permalink
Minor: Add more docs and examples for Transformed and `TransformedR…
Browse files Browse the repository at this point in the history
…esult` (apache#11003)
  • Loading branch information
alamb authored and xinlifoobar committed Jun 22, 2024
1 parent 07dceeb commit 0c3b068
Showing 1 changed file with 81 additions and 2 deletions.
83 changes: 81 additions & 2 deletions datafusion/common/src/tree_node.rs
Original file line number Diff line number Diff line change
Expand Up @@ -582,7 +582,11 @@ impl TreeNodeRecursion {

/// Result of tree walk / transformation APIs
///
/// API users control the transformation by returning:
/// `Transformed` is a wrapper around the tree node data (e.g. `Expr` or
/// `LogicalPlan`). It is used to indicate whether the node was transformed
/// and how the recursion should proceed.
///
/// [`TreeNode`] API users control the transformation by returning:
/// - The resulting (possibly transformed) node,
/// - `transformed`: flag indicating whether any change was made to the node
/// - `tnr`: [`TreeNodeRecursion`] specifying how to proceed with the recursion.
Expand All @@ -592,7 +596,66 @@ impl TreeNodeRecursion {
/// - `transformed`: flag indicating whether any change was made to the node
/// - `tnr`: [`TreeNodeRecursion`] specifying how the recursion ended.
///
/// Example APIs:
/// See also
/// * [`Transformed::update_data`] to modify the node without changing the `transformed` flag
/// * [`Transformed::map_data`] for fallable operation that return the same type
/// * [`Transformed::transform_data`] to chain fallable transformations
/// * [`TransformedResult`] for working with `Result<Transformed<U>>`
///
/// # Examples
///
/// Use [`Transformed::yes`] and [`Transformed::no`] to signal that a node was
/// rewritten and the recursion should continue:
///
/// ```
/// # use datafusion_common::tree_node::Transformed;
/// # // note use i64 instead of Expr as Expr is not in datafusion-common
/// # fn orig_expr() -> i64 { 1 }
/// # fn make_new_expr(i: i64) -> i64 { 2 }
/// let expr = orig_expr();
///
/// // Create a new `Transformed` object signaling the node was not rewritten
/// let ret = Transformed::no(expr.clone());
/// assert!(!ret.transformed);
///
/// // Create a new `Transformed` object signaling the node was rewritten
/// let ret = Transformed::yes(expr);
/// assert!(ret.transformed)
/// ```
///
/// Access the node within the `Transformed` object:
/// ```
/// # use datafusion_common::tree_node::Transformed;
/// # // note use i64 instead of Expr as Expr is not in datafusion-common
/// # fn orig_expr() -> i64 { 1 }
/// # fn make_new_expr(i: i64) -> i64 { 2 }
/// let expr = orig_expr();
///
/// // `Transformed` object signaling the node was not rewritten
/// let ret = Transformed::no(expr.clone());
/// // Access the inner object using .data
/// assert_eq!(expr, ret.data);
/// ```
///
/// Transform the node within the `Transformed` object.
///
/// ```
/// # use datafusion_common::tree_node::Transformed;
/// # // note use i64 instead of Expr as Expr is not in datafusion-common
/// # fn orig_expr() -> i64 { 1 }
/// # fn make_new_expr(i: i64) -> i64 { 2 }
/// let expr = orig_expr();
/// let ret = Transformed::no(expr.clone())
/// .transform_data(|expr| {
/// // closure returns a result and potentially transforms the node
/// // in this example, it does transform the node
/// let new_expr = make_new_expr(expr);
/// Ok(Transformed::yes(new_expr))
/// }).unwrap();
/// // transformed flag is the union of the original ans closure's transformed flag
/// assert!(ret.transformed);
/// ```
/// # Example APIs that use `TreeNode`
/// - [`TreeNode`],
/// - [`TreeNode::rewrite`],
/// - [`TreeNode::transform_down`],
Expand Down Expand Up @@ -833,6 +896,22 @@ macro_rules! map_until_stop_and_collect {
}

/// Transformation helper to access [`Transformed`] fields in a [`Result`] easily.
///
/// # Example
/// Access the internal data of a `Result<Transformed<T>>`
/// as a `Result<T>` using the `data` method:
/// ```
/// # use datafusion_common::Result;
/// # use datafusion_common::tree_node::{Transformed, TransformedResult};
/// # // note use i64 instead of Expr as Expr is not in datafusion-common
/// # fn update_expr() -> i64 { 1 }
/// # fn main() -> Result<()> {
/// let transformed: Result<Transformed<_>> = Ok(Transformed::yes(update_expr()));
/// // access the internal data of the transformed result, or return the error
/// let transformed_expr = transformed.data()?;
/// # Ok(())
/// # }
/// ```
pub trait TransformedResult<T> {
fn data(self) -> Result<T>;

Expand Down

0 comments on commit 0c3b068

Please sign in to comment.