Merge remote-tracking branch 'apache/main' into alamb/advanced_parque…

…t_index

datafusion-examples/examples/rewrite_expr.rs

@@ -18,13 +18,13 @@
 use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
-use datafusion_common::{plan_err, Result, ScalarValue};
+use datafusion_common::{plan_err, DataFusionError, Result, ScalarValue};
 use datafusion_expr::{
     AggregateUDF, Between, Expr, Filter, LogicalPlan, ScalarUDF, TableSource, WindowUDF,
 };
 use datafusion_optimizer::analyzer::{Analyzer, AnalyzerRule};
-use datafusion_optimizer::optimizer::Optimizer;
-use datafusion_optimizer::{utils, OptimizerConfig, OptimizerContext, OptimizerRule};
+use datafusion_optimizer::optimizer::{ApplyOrder, Optimizer};
+use datafusion_optimizer::{OptimizerConfig, OptimizerContext, OptimizerRule};
 use datafusion_sql::planner::{ContextProvider, SqlToRel};
 use datafusion_sql::sqlparser::dialect::PostgreSqlDialect;
 use datafusion_sql::sqlparser::parser::Parser;
@@ -133,30 +133,34 @@ impl OptimizerRule for MyOptimizerRule {
 
     fn try_optimize(
         &self,
-        plan: &LogicalPlan,
-        config: &dyn OptimizerConfig,
+        _plan: &LogicalPlan,
+        _config: &dyn OptimizerConfig,
     ) -> Result<Option<LogicalPlan>> {
-        // recurse down and optimize children first
-        let optimized_plan = utils::optimize_children(self, plan, config)?;
-        match optimized_plan {
-            Some(LogicalPlan::Filter(filter)) => {
+        unreachable!()
+    }
+
+    fn apply_order(&self) -> Option<ApplyOrder> {
+        Some(ApplyOrder::BottomUp)
+    }
+
+    fn supports_rewrite(&self) -> bool {
+        true
+    }
+
+    fn rewrite(
+        &self,
+        plan: LogicalPlan,
+        _config: &dyn OptimizerConfig,
+    ) -> Result<Transformed<LogicalPlan>, DataFusionError> {
+        match plan {
+            LogicalPlan::Filter(filter) => {
                 let predicate = my_rewrite(filter.predicate.clone())?;
-                Ok(Some(LogicalPlan::Filter(Filter::try_new(
+                Ok(Transformed::yes(LogicalPlan::Filter(Filter::try_new(
                     predicate,
-                    filter.input,
+                    filter.input.clone(),
                 )?)))
             }
-            Some(optimized_plan) => Ok(Some(optimized_plan)),
-            None => match plan {
-                LogicalPlan::Filter(filter) => {
-                    let predicate = my_rewrite(filter.predicate.clone())?;
-                    Ok(Some(LogicalPlan::Filter(Filter::try_new(
-                        predicate,
-                        filter.input.clone(),
-                    )?)))
-                }
-                _ => Ok(None),
-            },
+            _ => Ok(Transformed::no(plan)),
         }
     }
 }

datafusion/common/src/tree_node.rs

@@ -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.
@@ -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`],
@@ -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>;