buoyant_kernel 0.21.103

use std::borrow::{Cow, ToOwned};
use std::sync::Arc;

use crate::expressions::{
    BinaryExpression, BinaryPredicate, ColumnName, Expression, ExpressionRef, JunctionPredicate,
    MapToStructExpression, OpaqueExpression, OpaquePredicate, ParseJsonExpression, Predicate,
    Scalar, Transform, UnaryExpression, UnaryPredicate, VariadicExpression,
};
use crate::transforms::{map_owned_children_or_else, map_owned_or_else, map_owned_pair_or_else};

/// Generic framework for recursive bottom-up transforms of expressions and
/// predicates. Transformations return `Option<Cow>` with the following semantics:
///
/// * `Some(Cow::Owned)` -- The input was transformed and the parent should be updated with it.
/// * `Some(Cow::Borrowed)` -- The input was not transformed.
/// * `None` -- The input was filtered out and the parent should be updated to not reference it.
///
/// The transform entry point is generally [`Self::transform_expr`] or [`Self::transform_pred`] (for
/// expressions or predicates, respectively), but callers can also directly invoke the transform
/// for a specific expression/predicate variant (e.g. [`Self::transform_expr_column`] for
/// [`ColumnName`] or [`Self::transform_pred_unary`] for [`UnaryPredicate`]).
///
/// The provided `transform_xxx` methods all default to no-op (usually by invoking the corresponding
/// recursive helper method), and implementations should selectively override specific
/// `transform_xxx` methods as needed for the task at hand.
///
/// # Recursive helper methods
///
/// The provided `recurse_into_xxx` methods encapsulate the boilerplate work of recursing into the
/// child expression of each expression type. Except as specifically noted otherwise, these
/// recursive helpers all behave uniformly, based on the number of children the parent has:
///
/// * Leaf (no children) - Leaf `transform_xxx` methods simply return their argument unchanged, and
///   no corresponding `recurse_into_xxx` method is provided.
///
/// * Unary (single child) - If the child was filtered out, filter out the parent. If the child
///   changed, build a new parent around it. Otherwise, return the parent unchanged.
///
/// * Binary (two children) - If either child was filtered out, filter out the parent. If at least
///   one child changed, build a new parent around them. Otherwise, return the parent unchanged.
///
/// * Variadic (0+ children) - If no children remain (all filtered out), filter out the parent.
///   Otherwise, if at least one child changed or was filtered out, build a new parent around the
///   children. Otherwise, return the parent unchanged.
///
/// Implementations can call these as needed but will generally not need to override them.
pub trait ExpressionTransform<'a> {
    /// Called for each literal encountered during the traversal (leaf).
    fn transform_expr_literal(&mut self, value: &'a Scalar) -> Option<Cow<'a, Scalar>> {
        Some(Cow::Borrowed(value))
    }

    /// Called for each column reference encountered during the traversal (leaf).
    fn transform_expr_column(&mut self, name: &'a ColumnName) -> Option<Cow<'a, ColumnName>> {
        Some(Cow::Borrowed(name))
    }

    /// Called for the expression list of each struct expression encountered during the
    /// traversal. The provided implementation just forwards to [`Self::recurse_into_expr_struct`].
    fn transform_expr_struct(
        &mut self,
        fields: &'a [ExpressionRef],
    ) -> Option<Cow<'a, [ExpressionRef]>> {
        self.recurse_into_expr_struct(fields)
    }

    /// Called for each opaque expression encountered during the traversal. The provided
    /// implementation just forwards to [`Self::recurse_into_expr_opaque`].
    fn transform_expr_opaque(
        &mut self,
        expr: &'a OpaqueExpression,
    ) -> Option<Cow<'a, OpaqueExpression>> {
        self.recurse_into_expr_opaque(expr)
    }

    /// Called for each unknown expression encountered during the traversal (leaf).
    fn transform_expr_unknown(&mut self, name: &'a String) -> Option<Cow<'a, String>> {
        Some(Cow::Borrowed(name))
    }

    /// Called for each transform expression encountered during the traversal (leaf).
    ///
    /// The provided implementation does _NOT_ recurse into its children.
    fn transform_expr_transform(&mut self, transform: &'a Transform) -> Option<Cow<'a, Transform>> {
        Some(Cow::Borrowed(transform))
    }

    /// Called for each parse-json expression encountered during the traversal. The provided
    /// implementation just forwards to [`Self::recurse_into_expr_parse_json`].
    fn transform_expr_parse_json(
        &mut self,
        expr: &'a ParseJsonExpression,
    ) -> Option<Cow<'a, ParseJsonExpression>> {
        self.recurse_into_expr_parse_json(expr)
    }

    /// Called for each map-to-struct expression encountered during the traversal. The provided
    /// implementation just forwards to [`Self::recurse_into_expr_map_to_struct`].
    fn transform_expr_map_to_struct(
        &mut self,
        expr: &'a MapToStructExpression,
    ) -> Option<Cow<'a, MapToStructExpression>> {
        self.recurse_into_expr_map_to_struct(expr)
    }

    /// Called for the child of each predicate expression encountered during the
    /// traversal. The provided implementation just forwards to [`Self::recurse_into_expr_pred`].
    fn transform_expr_pred(&mut self, pred: &'a Predicate) -> Option<Cow<'a, Predicate>> {
        self.recurse_into_expr_pred(pred)
    }

    /// Called for the child of each NOT predicate encountered during the
    /// traversal. The provided implementation just forwards to [`Self::recurse_into_pred_not`].
    fn transform_pred_not(&mut self, pred: &'a Predicate) -> Option<Cow<'a, Predicate>> {
        self.recurse_into_pred_not(pred)
    }

    /// Called for each unary expression encountered during the traversal. The provided
    /// implementation just forwards to [`Self::recurse_into_expr_unary`].
    fn transform_expr_unary(
        &mut self,
        expr: &'a UnaryExpression,
    ) -> Option<Cow<'a, UnaryExpression>> {
        self.recurse_into_expr_unary(expr)
    }

    /// Called for each unary predicate encountered during the traversal. The provided
    /// implementation just forwards to [`Self::recurse_into_pred_unary`].
    fn transform_pred_unary(
        &mut self,
        pred: &'a UnaryPredicate,
    ) -> Option<Cow<'a, UnaryPredicate>> {
        self.recurse_into_pred_unary(pred)
    }

    /// Called for each binary expression encountered during the traversal. The provided
    /// implementation just forwards to [`Self::recurse_into_expr_binary`].
    fn transform_expr_binary(
        &mut self,
        expr: &'a BinaryExpression,
    ) -> Option<Cow<'a, BinaryExpression>> {
        self.recurse_into_expr_binary(expr)
    }

    /// Called for each binary predicate encountered during the traversal. The provided
    /// implementation just forwards to [`Self::recurse_into_pred_binary`].
    fn transform_pred_binary(
        &mut self,
        pred: &'a BinaryPredicate,
    ) -> Option<Cow<'a, BinaryPredicate>> {
        self.recurse_into_pred_binary(pred)
    }

    /// Called for each variadic expression encountered during the traversal. The provided
    /// implementation just forwards to [`Self::recurse_into_expr_variadic`].
    fn transform_expr_variadic(
        &mut self,
        expr: &'a VariadicExpression,
    ) -> Option<Cow<'a, VariadicExpression>> {
        self.recurse_into_expr_variadic(expr)
    }

    /// Called for each junction predicate encountered during the traversal. The provided
    /// implementation just forwards to [`Self::recurse_into_pred_junction`].
    fn transform_pred_junction(
        &mut self,
        pred: &'a JunctionPredicate,
    ) -> Option<Cow<'a, JunctionPredicate>> {
        self.recurse_into_pred_junction(pred)
    }

    /// Called for each opaque predicate encountered during the traversal. The provided
    /// implementation just forwards to [`Self::recurse_into_pred_opaque`].
    fn transform_pred_opaque(
        &mut self,
        pred: &'a OpaquePredicate,
    ) -> Option<Cow<'a, OpaquePredicate>> {
        self.recurse_into_pred_opaque(pred)
    }

    /// Called for each unknown predicate encountered during the traversal (leaf).
    fn transform_pred_unknown(&mut self, name: &'a String) -> Option<Cow<'a, String>> {
        Some(Cow::Borrowed(name))
    }

    /// General entry point for transforming an expression. This method will dispatch to the
    /// specific transform for each expression variant. Also invoked internally in order to recurse
    /// on the child(ren) of non-leaf expressions.
    fn transform_expr(&mut self, expr: &'a Expression) -> Option<Cow<'a, Expression>> {
        match expr {
            Expression::Literal(s) => {
                let child = self.transform_expr_literal(s);
                map_owned_or_else(expr, child, Expression::Literal)
            }
            Expression::Column(c) => {
                let child = self.transform_expr_column(c);
                map_owned_or_else(expr, child, Expression::Column)
            }
            Expression::Predicate(p) => {
                let child = self.transform_expr_pred(p);
                map_owned_or_else(expr, child, Expression::from)
            }
            Expression::Struct(s, nullability) => {
                let map_owned = |exprs| Expression::Struct(exprs, nullability.clone());
                map_owned_or_else(expr, self.transform_expr_struct(s), map_owned)
            }
            Expression::Transform(t) => {
                let child = self.transform_expr_transform(t);
                map_owned_or_else(expr, child, Expression::Transform)
            }
            Expression::Unary(u) => {
                let child = self.transform_expr_unary(u);
                map_owned_or_else(expr, child, Expression::Unary)
            }
            Expression::Binary(b) => {
                let child = self.transform_expr_binary(b);
                map_owned_or_else(expr, child, Expression::Binary)
            }
            Expression::Variadic(v) => {
                let child = self.transform_expr_variadic(v);
                map_owned_or_else(expr, child, Expression::Variadic)
            }
            Expression::Opaque(o) => {
                let child = self.transform_expr_opaque(o);
                map_owned_or_else(expr, child, Expression::Opaque)
            }
            Expression::ParseJson(p) => {
                let child = self.transform_expr_parse_json(p);
                map_owned_or_else(expr, child, Expression::ParseJson)
            }
            Expression::MapToStruct(m) => {
                let child = self.transform_expr_map_to_struct(m);
                map_owned_or_else(expr, child, Expression::MapToStruct)
            }
            Expression::Unknown(u) => {
                let child = self.transform_expr_unknown(u);
                map_owned_or_else(expr, child, Expression::Unknown)
            }
        }
    }

    /// General entry point for transforming a predicate. This method will dispatch to the specific
    /// transform for each predicate variant. Also invoked internally in order to recurse on the
    /// child(ren) of non-leaf variants.
    fn transform_pred(&mut self, pred: &'a Predicate) -> Option<Cow<'a, Predicate>> {
        match pred {
            Predicate::BooleanExpression(e) => {
                let child = self.transform_expr(e);
                map_owned_or_else(pred, child, Predicate::BooleanExpression)
            }
            Predicate::Not(p) => {
                let child = self.transform_pred_not(p);
                map_owned_or_else(pred, child, |p| p)
            }
            Predicate::Unary(u) => {
                let child = self.transform_pred_unary(u);
                map_owned_or_else(pred, child, Predicate::Unary)
            }
            Predicate::Binary(b) => {
                let child = self.transform_pred_binary(b);
                map_owned_or_else(pred, child, Predicate::Binary)
            }
            // Route through the constructor to normalize in case the transform removed children.
            // When `transform_pred` returns `None` for a child, it is filtered out, which may
            // reduce the junction to one or zero elements. The constructor normalizes these.
            Predicate::Junction(j) => {
                let child = self.transform_pred_junction(j);
                map_owned_or_else(pred, child, |j| Predicate::junction(j.op, j.preds))
            }
            Predicate::Opaque(o) => {
                let child = self.transform_pred_opaque(o);
                map_owned_or_else(pred, child, Predicate::Opaque)
            }
            Predicate::Unknown(u) => {
                let child = self.transform_pred_unknown(u);
                map_owned_or_else(pred, child, Predicate::Unknown)
            }
        }
    }

    /// Recursively transforms a struct's child expressions (variadic).
    fn recurse_into_expr_struct(
        &mut self,
        fields: &'a [ExpressionRef],
    ) -> Option<Cow<'a, [ExpressionRef]>> {
        let children = fields.iter().map(|f| -> Option<Cow<'a, ExpressionRef>> {
            map_owned_or_else(f, self.transform_expr(f), Arc::new)
        });
        map_owned_children_or_else(fields, children, |fields| fields)
    }

    /// Recursively transforms the child expression of a parse-json expression (unary).
    fn recurse_into_expr_parse_json(
        &mut self,
        expr: &'a ParseJsonExpression,
    ) -> Option<Cow<'a, ParseJsonExpression>> {
        let f = |json_expr| ParseJsonExpression::new(json_expr, expr.output_schema.clone());
        map_owned_or_else(expr, self.transform_expr(&expr.json_expr), f)
    }

    /// Recursively transforms the child expression of a map-to-struct expression (unary).
    fn recurse_into_expr_map_to_struct(
        &mut self,
        expr: &'a MapToStructExpression,
    ) -> Option<Cow<'a, MapToStructExpression>> {
        let nested = self.transform_expr(&expr.map_expr);
        map_owned_or_else(expr, nested, MapToStructExpression::new)
    }

    /// Recursively transforms the children of an opaque expression (variadic).
    fn recurse_into_expr_opaque(
        &mut self,
        o: &'a OpaqueExpression,
    ) -> Option<Cow<'a, OpaqueExpression>> {
        let transformed_children = o.exprs.iter().map(|e| self.transform_expr(e));
        let map_owned = |exprs| OpaqueExpression::new(o.op.clone(), exprs);
        map_owned_children_or_else(o, transformed_children, map_owned)
    }

    /// Recursively transforms the child of a predicate expression (unary).
    fn recurse_into_expr_pred(&mut self, pred: &'a Predicate) -> Option<Cow<'a, Predicate>> {
        self.transform_pred(pred)
    }

    /// Recursively transforms the child of a not predicate expression (unary).
    fn recurse_into_pred_not(&mut self, p: &'a Predicate) -> Option<Cow<'a, Predicate>> {
        map_owned_or_else(p, self.transform_pred(p), Predicate::not)
    }

    /// Recursively transforms a unary predicate's child (unary).
    fn recurse_into_pred_unary(
        &mut self,
        u: &'a UnaryPredicate,
    ) -> Option<Cow<'a, UnaryPredicate>> {
        let nested = self.transform_expr(&u.expr);
        map_owned_or_else(u, nested, |expr| UnaryPredicate::new(u.op, expr))
    }

    /// Recursively transforms a binary predicate's children (binary).
    fn recurse_into_pred_binary(
        &mut self,
        b: &'a BinaryPredicate,
    ) -> Option<Cow<'a, BinaryPredicate>> {
        let left = self.transform_expr(&b.left);
        let right = self.transform_expr(&b.right);
        let f = |(left, right)| BinaryPredicate::new(b.op, left, right);
        map_owned_pair_or_else(b, left, right, f)
    }

    /// Recursively transforms a unary expression's child (unary).
    fn recurse_into_expr_unary(
        &mut self,
        u: &'a UnaryExpression,
    ) -> Option<Cow<'a, UnaryExpression>> {
        let nested = self.transform_expr(&u.expr);
        map_owned_or_else(u, nested, |expr| UnaryExpression::new(u.op, expr))
    }

    /// Recursively transforms a binary expression's children (binary).
    fn recurse_into_expr_binary(
        &mut self,
        b: &'a BinaryExpression,
    ) -> Option<Cow<'a, BinaryExpression>> {
        let left = self.transform_expr(&b.left);
        let right = self.transform_expr(&b.right);
        let f = |(left, right)| BinaryExpression::new(b.op, left, right);
        map_owned_pair_or_else(b, left, right, f)
    }

    /// Recursively transforms a variadic expression's children (variadic).
    fn recurse_into_expr_variadic(
        &mut self,
        v: &'a VariadicExpression,
    ) -> Option<Cow<'a, VariadicExpression>> {
        let children = v.exprs.iter().map(|e| self.transform_expr(e));
        map_owned_children_or_else(v, children, |exprs| VariadicExpression::new(v.op, exprs))
    }

    /// Recursively transforms a junction predicate's children (variadic).
    fn recurse_into_pred_junction(
        &mut self,
        j: &'a JunctionPredicate,
    ) -> Option<Cow<'a, JunctionPredicate>> {
        let children = j.preds.iter().map(|p| self.transform_pred(p));
        map_owned_children_or_else(j, children, |preds| JunctionPredicate::new(j.op, preds))
    }

    /// Recursively transforms an opaque predicate's children (variadic).
    fn recurse_into_pred_opaque(
        &mut self,
        o: &'a OpaquePredicate,
    ) -> Option<Cow<'a, OpaquePredicate>> {
        let children = o.exprs.iter().map(|e| self.transform_expr(e));
        let map_owned = |exprs| OpaquePredicate::new(o.op.clone(), exprs);
        map_owned_children_or_else(o, children, map_owned)
    }
}

/// An expression "transform" that doesn't actually change the expression at all. Instead, it
/// measures the maximum depth of a expression, with a depth limit to prevent stack overflow. Useful
/// for verifying that a expression has reasonable depth before attempting to work with it.
pub struct ExpressionDepthChecker {
    depth_limit: usize,
    max_depth_seen: usize,
    current_depth: usize,
    call_count: usize,
}

impl ExpressionDepthChecker {
    /// Depth-checks the given expression against a given depth limit. The return value is the
    /// largest depth seen, which is capped at one more than the depth limit (indicating the
    /// recursion was terminated).
    pub fn check_expr(expr: &Expression, depth_limit: usize) -> usize {
        Self::check_expr_with_call_count(expr, depth_limit).0
    }

    /// Depth-checks the given predicate against a given depth limit. The return value is the
    /// largest depth seen, which is capped at one more than the depth limit (indicating the
    /// recursion was terminated).
    pub fn check_pred(pred: &Predicate, depth_limit: usize) -> usize {
        Self::check_pred_with_call_count(pred, depth_limit).0
    }

    // Exposed for testing
    fn check_expr_with_call_count(expr: &Expression, depth_limit: usize) -> (usize, usize) {
        let mut checker = Self::new(depth_limit);
        let _ = checker.transform_expr(expr);
        (checker.max_depth_seen, checker.call_count)
    }

    // Exposed for testing
    fn check_pred_with_call_count(pred: &Predicate, depth_limit: usize) -> (usize, usize) {
        let mut checker = Self::new(depth_limit);
        let _ = checker.transform_pred(pred);
        (checker.max_depth_seen, checker.call_count)
    }

    fn new(depth_limit: usize) -> Self {
        Self {
            depth_limit,
            max_depth_seen: 0,
            current_depth: 0,
            call_count: 0,
        }
    }

    // Triggers the requested recursion only doing so would not exceed the depth limit.
    fn depth_limited<'a, T: std::fmt::Debug + ToOwned + ?Sized>(
        &mut self,
        recurse: impl FnOnce(&mut Self, &'a T) -> Option<Cow<'a, T>>,
        arg: &'a T,
    ) -> Option<Cow<'a, T>> {
        self.call_count += 1;
        if self.max_depth_seen < self.current_depth {
            self.max_depth_seen = self.current_depth;
            if self.depth_limit < self.current_depth {
                tracing::warn!(
                    "Max expression depth {} exceeded by {arg:?}",
                    self.depth_limit
                );
            }
        }
        if self.max_depth_seen <= self.depth_limit {
            self.current_depth += 1;
            let _ = recurse(self, arg);
            self.current_depth -= 1;
        }
        None
    }
}

impl<'a> ExpressionTransform<'a> for ExpressionDepthChecker {
    fn transform_expr_struct(
        &mut self,
        fields: &'a [ExpressionRef],
    ) -> Option<Cow<'a, [ExpressionRef]>> {
        self.depth_limited(Self::recurse_into_expr_struct, fields)
    }

    fn transform_expr_pred(&mut self, pred: &'a Predicate) -> Option<Cow<'a, Predicate>> {
        self.depth_limited(Self::recurse_into_expr_pred, pred)
    }

    fn transform_pred_not(&mut self, pred: &'a Predicate) -> Option<Cow<'a, Predicate>> {
        self.depth_limited(Self::recurse_into_pred_not, pred)
    }

    fn transform_pred_unary(
        &mut self,
        pred: &'a UnaryPredicate,
    ) -> Option<Cow<'a, UnaryPredicate>> {
        self.depth_limited(Self::recurse_into_pred_unary, pred)
    }

    fn transform_expr_binary(
        &mut self,
        expr: &'a BinaryExpression,
    ) -> Option<Cow<'a, BinaryExpression>> {
        self.depth_limited(Self::recurse_into_expr_binary, expr)
    }

    fn transform_pred_binary(
        &mut self,
        pred: &'a BinaryPredicate,
    ) -> Option<Cow<'a, BinaryPredicate>> {
        self.depth_limited(Self::recurse_into_pred_binary, pred)
    }

    fn transform_pred_junction(
        &mut self,
        pred: &'a JunctionPredicate,
    ) -> Option<Cow<'a, JunctionPredicate>> {
        self.depth_limited(Self::recurse_into_pred_junction, pred)
    }

    fn transform_pred_opaque(
        &mut self,
        pred: &'a OpaquePredicate,
    ) -> Option<Cow<'a, OpaquePredicate>> {
        self.depth_limited(Self::recurse_into_pred_opaque, pred)
    }

    fn transform_expr_opaque(
        &mut self,
        expr: &'a OpaqueExpression,
    ) -> Option<Cow<'a, OpaqueExpression>> {
        self.depth_limited(Self::recurse_into_expr_opaque, expr)
    }

    fn transform_expr_map_to_struct(
        &mut self,
        expr: &'a MapToStructExpression,
    ) -> Option<Cow<'a, MapToStructExpression>> {
        self.depth_limited(Self::recurse_into_expr_map_to_struct, expr)
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use super::*;
    use crate::expressions::VariadicExpressionOp::Coalesce;
    use crate::expressions::{
        column_expr, column_pred, Expression, Expression as Expr, OpaqueExpressionOp,
        OpaquePredicateOp, ParseJsonExpression, Predicate as Pred, Scalar,
        ScalarExpressionEvaluator, VariadicExpression,
    };
    use crate::kernel_predicates::{
        DirectDataSkippingPredicateEvaluator, DirectPredicateEvaluator,
        IndirectDataSkippingPredicateEvaluator,
    };
    use crate::schema::{DataType, StructField, StructType};
    use crate::DeltaResult;

    #[derive(Debug, PartialEq)]
    struct OpaqueTestOp(String);

    impl OpaqueExpressionOp for OpaqueTestOp {
        fn name(&self) -> &str {
            &self.0
        }
        fn eval_expr_scalar(
            &self,
            _eval_expr: &ScalarExpressionEvaluator<'_>,
            _exprs: &[Expression],
        ) -> DeltaResult<Scalar> {
            unimplemented!()
        }
    }

    impl OpaquePredicateOp for OpaqueTestOp {
        fn name(&self) -> &str {
            &self.0
        }

        fn eval_pred_scalar(
            &self,
            _eval_expr: &ScalarExpressionEvaluator<'_>,
            _evaluator: &DirectPredicateEvaluator<'_>,
            _exprs: &[Expr],
            _inverted: bool,
        ) -> DeltaResult<Option<bool>> {
            unimplemented!()
        }

        fn eval_as_data_skipping_predicate(
            &self,
            _predicate_evaluator: &DirectDataSkippingPredicateEvaluator<'_>,
            _exprs: &[Expr],
            _inverted: bool,
        ) -> Option<bool> {
            unimplemented!()
        }

        fn as_data_skipping_predicate(
            &self,
            _predicate_evaluator: &IndirectDataSkippingPredicateEvaluator<'_>,
            _exprs: &[Expr],
            _inverted: bool,
        ) -> Option<Pred> {
            unimplemented!()
        }
    }

    struct NoopTransform;
    impl ExpressionTransform<'_> for NoopTransform {}

    struct ColumnReplacer;
    impl<'a> ExpressionTransform<'a> for ColumnReplacer {
        fn transform_expr_column(&mut self, name: &'a ColumnName) -> Option<Cow<'a, ColumnName>> {
            if name.len() == 1 && name[0] == "old_col" {
                Some(Cow::Owned(ColumnName::new(["new_col"])))
            } else {
                Some(Cow::Borrowed(name))
            }
        }
    }

    #[test]
    fn test_transform_expr_variadic_noop() {
        // Test default no-op behavior - should return Cow::Borrowed
        let variadic_expr = VariadicExpression::new(
            Coalesce,
            vec![Expr::literal(1), column_expr!("x"), Expr::literal("test")],
        );

        let mut transform = NoopTransform;
        let result = transform.transform_expr_variadic(&variadic_expr);

        assert!(matches!(result, Some(Cow::Borrowed(_))));
        if let Some(Cow::Borrowed(result_expr)) = result {
            assert_eq!(result_expr, &variadic_expr);
        }
    }

    #[test]
    fn test_transform_expr_variadic_empty_input() {
        // Test edge case with empty children list
        let variadic_expr = VariadicExpression::new(Coalesce, Vec::<Expr>::new());

        let mut transform = NoopTransform;
        let result = transform.transform_expr_variadic(&variadic_expr);

        // Empty children list with no-op transform returns None because new_children.is_empty()
        // This is the behavior of recurse_into_children when starting with empty slice
        assert!(result.is_none());
    }

    #[test]
    fn test_transform_expr_variadic_child_transformation() {
        // Test transformation of child expressions - should return Cow::Owned
        let variadic_expr = VariadicExpression::new(
            Coalesce,
            vec![
                Expr::literal(1),
                column_expr!("old_col"),
                column_expr!("unchanged_col"),
                Expr::literal("test"),
            ],
        );

        let result = ColumnReplacer.transform_expr_variadic(&variadic_expr);

        assert!(matches!(result, Some(Cow::Owned(_))));
        if let Some(Cow::Owned(result_expr)) = result {
            assert_eq!(result_expr.op, Coalesce);
            assert_eq!(result_expr.exprs.len(), 4);

            // Check that the column was replaced
            if let Expr::Column(col) = &result_expr.exprs[1] {
                assert_eq!(col.len(), 1);
                assert_eq!(col[0], "new_col");
            } else {
                panic!("Expected column expression");
            }

            // Check that other expressions are unchanged
            assert_eq!(result_expr.exprs[0], Expr::literal(1));
            if let Expr::Column(col) = &result_expr.exprs[2] {
                assert_eq!(col.len(), 1);
                assert_eq!(col[0], "unchanged_col");
            } else {
                panic!("Expected column expression");
            }
            assert_eq!(result_expr.exprs[3], Expr::literal("test"));
        }
    }

    #[test]
    fn test_transform_expr_variadic_child_removal() {
        // Test removal of child expressions - should return Cow::Owned with fewer children
        struct LiteralRemover;
        impl<'a> ExpressionTransform<'a> for LiteralRemover {
            fn transform_expr_literal(&mut self, _value: &'a Scalar) -> Option<Cow<'a, Scalar>> {
                None // Remove all literals
            }
        }

        let variadic_expr = VariadicExpression::new(
            Coalesce,
            vec![
                Expr::literal(1),
                column_expr!("x"),
                Expr::literal("test"),
                column_expr!("y"),
            ],
        );

        let mut transform = LiteralRemover;
        let result = transform.transform_expr_variadic(&variadic_expr);

        assert!(matches!(result, Some(Cow::Owned(_))));
        if let Some(Cow::Owned(result_expr)) = result {
            assert_eq!(result_expr.op, Coalesce);
            assert_eq!(result_expr.exprs.len(), 2); // Only columns should remain

            // Check that only columns remain
            if let Expr::Column(col) = &result_expr.exprs[0] {
                assert_eq!(col.len(), 1);
                assert_eq!(col[0], "x");
            } else {
                panic!("Expected column expression");
            }
            if let Expr::Column(col) = &result_expr.exprs[1] {
                assert_eq!(col.len(), 1);
                assert_eq!(col[0], "y");
            } else {
                panic!("Expected column expression");
            }
        }
    }

    #[test]
    fn test_transform_expr_variadic_all_children_removed() {
        // Test edge case where all children are removed - should return None
        struct RemoveAll;
        impl<'a> ExpressionTransform<'a> for RemoveAll {
            fn transform_expr_literal(&mut self, _value: &'a Scalar) -> Option<Cow<'a, Scalar>> {
                None
            }
            fn transform_expr_column(
                &mut self,
                _name: &'a ColumnName,
            ) -> Option<Cow<'a, ColumnName>> {
                None
            }
        }

        let variadic_expr = VariadicExpression::new(
            Coalesce,
            vec![Expr::literal(1), column_expr!("x"), Expr::literal("test")],
        );

        let mut transform = RemoveAll;
        let result = transform.transform_expr_variadic(&variadic_expr);

        assert!(result.is_none());
    }

    #[test]
    fn test_transform_expr_variadic_mixed_transformations() {
        // Test mixed scenario: some children transformed, some removed, some unchanged
        struct MixedTransform;
        impl<'a> ExpressionTransform<'a> for MixedTransform {
            fn transform_expr_literal(&mut self, value: &'a Scalar) -> Option<Cow<'a, Scalar>> {
                match value {
                    Scalar::Integer(1) => None,                 // Remove literal 1
                    Scalar::String(s) if s == "remove" => None, // Remove "remove" string
                    Scalar::Integer(n) => Some(Cow::Owned(Scalar::Integer(n * 2))), /* Double other integers */
                    _ => Some(Cow::Borrowed(value)), // Keep others unchanged
                }
            }
            fn transform_expr_column(
                &mut self,
                name: &'a ColumnName,
            ) -> Option<Cow<'a, ColumnName>> {
                if name.len() == 1 && name[0] == "transform_me" {
                    Some(Cow::Owned(ColumnName::new(["transformed"])))
                } else {
                    Some(Cow::Borrowed(name))
                }
            }
        }

        let variadic_expr = VariadicExpression::new(
            Coalesce,
            vec![
                Expr::literal(1),             // Will be removed
                column_expr!("unchanged"),    // Will stay unchanged
                Expr::literal(5),             // Will be transformed to 10
                Expr::literal("remove"),      // Will be removed
                column_expr!("transform_me"), // Will be transformed
                Expr::literal("keep"),        // Will stay unchanged
            ],
        );

        let mut transform = MixedTransform;
        let result = transform.transform_expr_variadic(&variadic_expr);

        assert!(matches!(result, Some(Cow::Owned(_))));
        if let Some(Cow::Owned(result_expr)) = result {
            assert_eq!(result_expr.op, Coalesce);
            assert_eq!(result_expr.exprs.len(), 4); // 2 removed, 4 remaining

            // Check remaining expressions in order
            if let Expr::Column(col) = &result_expr.exprs[0] {
                assert_eq!(col.len(), 1);
                assert_eq!(col[0], "unchanged");
            } else {
                panic!("Expected unchanged column");
            }

            assert_eq!(result_expr.exprs[1], Expr::literal(10)); // 5 * 2

            if let Expr::Column(col) = &result_expr.exprs[2] {
                assert_eq!(col.len(), 1);
                assert_eq!(col[0], "transformed");
            } else {
                panic!("Expected transformed column");
            }

            assert_eq!(result_expr.exprs[3], Expr::literal("keep"));
        }
    }

    fn test_output_schema() -> Arc<StructType> {
        Arc::new(StructType::new_unchecked(vec![
            StructField::new("a", DataType::LONG, true),
            StructField::new("b", DataType::STRING, true),
        ]))
    }

    #[test]
    fn test_transform_expr_parse_json_noop() {
        // Test default no-op behavior - should return Cow::Borrowed
        let parse_json_expr =
            ParseJsonExpression::new(column_expr!("json_col"), test_output_schema());

        let mut transform = NoopTransform;
        let result = transform.transform_expr_parse_json(&parse_json_expr);

        assert!(matches!(result, Some(Cow::Borrowed(_))));
        if let Some(Cow::Borrowed(result_expr)) = result {
            assert_eq!(result_expr, &parse_json_expr);
        }
    }

    #[test]
    fn test_transform_expr_parse_json_child_transformation() {
        // Test transformation of child expression - should return Cow::Owned
        let parse_json_expr =
            ParseJsonExpression::new(column_expr!("old_col"), test_output_schema());

        let result = ColumnReplacer.transform_expr_parse_json(&parse_json_expr);

        assert!(matches!(result, Some(Cow::Owned(_))));
        if let Some(Cow::Owned(result_expr)) = result {
            // Check that the column was replaced
            if let Expr::Column(col) = result_expr.json_expr.as_ref() {
                assert_eq!(col.len(), 1);
                assert_eq!(col[0], "new_col");
            } else {
                panic!("Expected column expression");
            }
            // Schema should be preserved
            assert_eq!(result_expr.output_schema, test_output_schema());
        }
    }

    #[test]
    fn test_transform_expr_parse_json_child_unchanged() {
        // Test when child column doesn't match replacement criteria - should return Cow::Borrowed
        let parse_json_expr =
            ParseJsonExpression::new(column_expr!("unchanged_col"), test_output_schema());

        let result = ColumnReplacer.transform_expr_parse_json(&parse_json_expr);

        // Since "unchanged_col" doesn't match "old_col", nothing changes
        assert!(matches!(result, Some(Cow::Borrowed(_))));
    }

    #[test]
    fn test_transform_expr_parse_json_child_removal() {
        // Test removal of child expression - should return None
        struct ColumnRemover;
        impl<'a> ExpressionTransform<'a> for ColumnRemover {
            fn transform_expr_column(
                &mut self,
                _name: &'a ColumnName,
            ) -> Option<Cow<'a, ColumnName>> {
                None // Remove all column references
            }
        }

        let parse_json_expr =
            ParseJsonExpression::new(column_expr!("json_col"), test_output_schema());

        let mut transform = ColumnRemover;
        let result = transform.transform_expr_parse_json(&parse_json_expr);

        // Child was removed, so the whole ParseJson should be None
        assert!(result.is_none());
    }

    #[test]
    fn test_transform_expr_parse_json_nested_child() {
        // Test with a more complex nested child expression
        struct LiteralDoubler;
        impl<'a> ExpressionTransform<'a> for LiteralDoubler {
            fn transform_expr_literal(&mut self, value: &'a Scalar) -> Option<Cow<'a, Scalar>> {
                if let Scalar::Integer(n) = value {
                    Some(Cow::Owned(Scalar::Integer(n * 2)))
                } else {
                    Some(Cow::Borrowed(value))
                }
            }
        }

        // ParseJson with a binary expression as child: column + 5
        let child_expr = column_expr!("x") + Expr::literal(5);
        let parse_json_expr = ParseJsonExpression::new(child_expr, test_output_schema());

        let mut transform = LiteralDoubler;
        let result = transform.transform_expr_parse_json(&parse_json_expr);

        assert!(matches!(result, Some(Cow::Owned(_))));
        if let Some(Cow::Owned(result_expr)) = result {
            // The literal 5 should have been doubled to 10
            if let Expr::Binary(binary) = result_expr.json_expr.as_ref() {
                if let Expr::Literal(Scalar::Integer(n)) = &*binary.right {
                    assert_eq!(*n, 10);
                } else {
                    panic!("Expected integer literal");
                }
            } else {
                panic!("Expected binary expression");
            }
        }
    }

    #[test]
    fn test_depth_checker() {
        let pred = Pred::or_from([
            Pred::and_from([
                Pred::opaque(
                    OpaqueTestOp("opaque".to_string()),
                    vec![
                        Expr::literal(10) + column_expr!("x"),
                        Expr::unknown("unknown") - column_expr!("b"),
                    ],
                ),
                Pred::literal(true),
                Pred::not(Pred::literal(true)),
            ]),
            Pred::and_from([
                Pred::is_null(column_expr!("b")),
                Pred::gt(Expr::literal(10), column_expr!("x")),
                Pred::or(
                    Pred::gt(
                        Expr::literal(5)
                            + Expr::opaque(
                                OpaqueTestOp("inscrutable".to_string()),
                                vec![Expr::literal(10)],
                            ),
                        Expr::literal(20),
                    ),
                    column_pred!("y"),
                ),
                Pred::unknown("mystery"),
            ]),
            Pred::eq(
                Expr::literal(42),
                Expr::struct_from([Expr::literal(10), column_expr!("b")]),
            ),
        ]);

        // Verify the default/no-op transform, since we have this nice complex expression handy.
        assert!(matches!(
            NoopTransform.transform_pred(&pred),
            Some(std::borrow::Cow::Borrowed(_))
        ));

        // Similar to ExpressionDepthChecker::check_pred, but also returns call count
        let check_with_call_count =
            |depth_limit| ExpressionDepthChecker::check_pred_with_call_count(&pred, depth_limit);

        // NOTE: The checker ignores leaf nodes!

        // OR
        //  * AND
        //    * OPAQUE   >LIMIT<
        //    * NOT
        //  * AND
        //  * EQ
        assert_eq!(check_with_call_count(1), (2, 6));

        // OR
        //  * AND
        //    * OPAQUE
        //      * PLUS      >LIMIT<
        //      * MINUS
        //    * NOT
        //  * AND
        //  * EQ
        assert_eq!(check_with_call_count(2), (3, 8));

        // OR
        //  * AND
        //    * OPAQUE
        //      * PLUS
        //      * MINUS
        //    * NOT
        //  * AND
        //    * IS NULL
        //    * GT
        //    * OR
        //      * GT
        //        * PLUS     >LIMIT<
        //  * EQ
        assert_eq!(check_with_call_count(3), (4, 13));

        // OR
        //  * AND
        //    * OPAQUE
        //      * PLUS
        //      * MINUS
        //    * NOT
        //  * AND
        //    * IS_NULL
        //    * GT
        //    * OR
        //      * GT
        //        * PLUS
        //          * OPAQUE    >LIMIT<
        //  * EQ
        assert_eq!(check_with_call_count(4), (5, 14));

        // Depth limit not hit (full traversal required)
        //
        // OR
        //  * AND
        //    * OPAQUE
        //      * PLUS
        //      * MINUS
        //    * NOT
        //  * AND
        //    * IS_NULL
        //    * GT
        //    * OR
        //      * GT
        //        * PLUS
        //          * OPAQUE
        //  * EQ
        //    * STRUCT
        assert_eq!(check_with_call_count(5), (5, 15));
        assert_eq!(check_with_call_count(6), (5, 15));

        // Check expressions as well
        let expr = Expr::from(pred);
        let check_with_call_count =
            |depth_limit| ExpressionDepthChecker::check_expr_with_call_count(&expr, depth_limit);

        // Adding an `Expression::Predicate` root makes the expression tree exactly one node taller,
        // which makes the recursion terminate sooner than previously:
        //
        // PRED
        //  * OR
        //    * AND              > LIMIT 1 <
        //      * OPAQUE         > LIMIT 2 <
        //        * PLUS         > LIMIT 3 <
        //        * MINUS
        //      * NOT
        //    * AND
        //      * IS_NULL
        //      * GT
        //      * OR
        //        * GT
        //          * PLUS       > LIMIT 4 <
        //            * OPAQUE   > LIMIT 5 <
        //    * EQ
        //      * STRUCT
        assert_eq!(check_with_call_count(1), (2, 5));
        assert_eq!(check_with_call_count(2), (3, 7));
        assert_eq!(check_with_call_count(3), (4, 9));
        assert_eq!(check_with_call_count(4), (5, 14));
        assert_eq!(check_with_call_count(5), (6, 15));
        assert_eq!(check_with_call_count(6), (6, 16));
        assert_eq!(check_with_call_count(7), (6, 16));
    }

    #[test]
    fn transform_junction_to_single_child_unwraps() {
        // A transform that removes one child from AND(a, b) should produce the surviving
        // predicate directly, not a degenerate single-element junction AND(a).
        struct LiteralRemover;
        impl<'a> ExpressionTransform<'a> for LiteralRemover {
            fn transform_expr_literal(&mut self, _value: &'a Scalar) -> Option<Cow<'a, Scalar>> {
                None
            }
        }

        let pred = Pred::and(column_pred!("x"), Pred::literal(true));
        let mut transform = LiteralRemover;
        let result = transform.transform_pred(&pred);
        let result = result.map(Cow::into_owned);
        assert_eq!(result.as_ref(), Some(&column_pred!("x")));
        assert!(!matches!(result, Some(Pred::Junction(_))));
    }

    #[test]
    fn transform_junction_removing_all_children_returns_none() {
        // Removing all children propagates None (the junction is dropped entirely),
        // rather than producing an empty junction or identity literal.
        struct ColumnRemover;
        impl<'a> ExpressionTransform<'a> for ColumnRemover {
            fn transform_expr_column(
                &mut self,
                _name: &'a ColumnName,
            ) -> Option<Cow<'a, ColumnName>> {
                None
            }
        }

        let pred = Pred::and(column_pred!("x"), column_pred!("y"));
        let mut transform = ColumnRemover;
        assert!(transform.transform_pred(&pred).is_none());
    }
}