adriangb commented on code in PR #25458: URL: https://github.com/apache/datafusion/pull/25458#discussion_r4420628343
########## datafusion/optimizer/src/evaluation_sites.rs: ########## @@ -0,0 +1,444 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership +// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements + +//! The "no new volatile evaluations" optimizer invariant. +//! +//! Several optimizer rules move a plan node below another node that *defines* +//! one of the columns it uses. To do that the rule inlines the definition of +//! that column. If the definition is volatile, the two copies give two +//! different answers and the query returns wrong results. Guards against this +//! exist per rule and per call site, and each one was added after a bug +//! report: <https://github.com/apache/datafusion/issues/24678>, +//! <https://github.com/apache/datafusion/issues/25415>, +//! <https://github.com/apache/datafusion/issues/23655>. +//! +//! This module states the same property one level up, as an invariant of the +//! whole optimizer run. +//! +//! # An end of run invariant, by design +//! +//! [`check_no_new_volatile_evaluations`] compares the final plan of an +//! optimizer run against the plan the optimizer was given. Tests call it, see +//! `datafusion/core/tests/sql/volatile_evaluations.rs`. Intermediate plans are +//! not checked, because +//! rules cooperate: `simplify_expressions` expands `x BETWEEN a AND b` into +//! two references to `x`, and `CommonSubexprEliminate` hoists the duplicate +//! back out in the same pass. Only the plan that the user pays for matters, +//! so only that plan is checked. +//! +//! The price is that the check does not name the rule at fault. The before +//! and after site counts in the message are enough to start from. +//! +//! # What is counted +//! +//! A *site* is one volatile function call: an [`Expr`] for which +//! [`Expr::is_volatile_node`] holds, so a scalar, aggregate, window or higher +//! order function. Volatile calls only: two evaluations of a volatile call +//! give two different answers, so a new site is always a wrong results bug. +//! +//! A new evaluation site of a call that is not volatile, such as an +//! [`ExpressionPlacement::KeepInPlace`](datafusion_expr_common::placement::ExpressionPlacement::KeepInPlace) +//! one, is not counted, because it is not an invariant: pushing a predicate +//! through a projection that computes such an expression one time is a decided +//! trade-off, since the predicate can then reach the scan. See the decisions +//! table of <https://github.com/apache/datafusion/issues/25459>. +//! +//! Sites are keyed by a structural form of the call, in which every column, +//! outer reference columns included, is reduced to its unqualified name. A rule that duplicates a call and +//! re-qualifies its columns, for example from `f(a)` to `f(test.a)`, must not +//! escape the check. Only a call that is already in the input plan is checked. +//! A rule is free to rewrite `concat_ws` into `concat`, or one Spark function +//! into two DataFusion functions. That is a new call, not a second evaluation +//! of an old one. +//! +//! # Row lineage +//! +//! A whole plan count has false positives: a predicate is legitimately +//! duplicated across the branches of a `Union` and across the two sides of a +//! join. So the count is taken per *row lineage*: the number of times a call +//! is evaluated on the way from one leaf to the root. +//! +//! ```text +//! sites(node) = own_sites(node) + max over inputs i of sites(i) +//! ``` +//! +//! The maximum, not the sum, because a row that enters the left input of a +//! join never meets the nodes of the right input. + +use std::collections::HashMap; +use std::sync::Arc; + +use datafusion_common::{ + Column, Result, internal_err, + tree_node::{Transformed, TransformedResult, TreeNode, TreeNodeRecursion}, +}; +use datafusion_expr::expr::{Exists, InSubquery, SetComparison}; +use datafusion_expr::{Expr, LogicalPlan}; + +/// The number of evaluation sites of one call, per row lineage, keyed by +/// [`structural_key`]. +type Counts = HashMap<String, usize>; + +/// The volatile evaluation sites of a [`LogicalPlan`]. +/// +/// See the [module documentation](self) for the definition. +#[derive(Debug, Default)] +pub(crate) struct EvaluationSites { + counts: Counts, +} + +impl EvaluationSites { + /// Count the evaluation sites of `plan`. + fn count(plan: &LogicalPlan) -> Self { + let mut counts = Counts::new(); + count_plan(plan, &mut counts); + Self { counts } + } + + /// Fail if `plan` evaluates a volatile call more times than `self` does. + /// + /// `self` must be the count of the plan that the optimizer was given. + fn check(&self, plan: &LogicalPlan) -> Result<()> { + let after = Self::count(plan); + let mut offenders: Vec<(&String, usize, usize)> = after + .counts + .iter() + .filter_map(|(key, count)| { + // A call that the optimizer invented is not a second + // evaluation of a call that was already there. + let before = *self.counts.get(key)?; + (*count > before).then_some((key, before, *count)) + }) + .collect(); + if offenders.is_empty() { + return Ok(()); + } + // Sorted so that the message does not depend on the hash order. + offenders.sort(); + + let details = offenders + .iter() + .map(|(call, before, after)| { + format!( + "volatile expression `{call}` is evaluated at {after} sites, \ + against {before} in the input plan" + ) + }) + .collect::<Vec<_>>() + .join("; "); + internal_err!( + "Optimizer added volatile evaluation sites: {details}. Two \ + evaluations of a volatile call give two different answers, so the \ + optimized plan must not evaluate one more times than the plan it \ + was given does. Counts are per row lineage, see `EvaluationSites`" + ) + } +} + +/// Fail if `optimized` evaluates a volatile call at more sites than `input` +/// does. `input` is the plan the optimizer was given. +/// +/// See the [module documentation](self) for what is counted. +pub fn check_no_new_volatile_evaluations( Review Comment: Do we have to make this pub? Can it be behind a `test` feature flag? -- 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]
