Skip to content

Commit ddb2155

Browse files
adriangbclaude
andcommitted
test: check that the logical optimizer reaches a fixed point
Record the plan after each optimizer pass and require the last pass to leave it unchanged. Covers the `floor_preimage` shape, the issue #25329 shape, and every statement of the 99 TPC-DS queries. Co-Authored-By: Claude Opus 5.5 <[email protected]>
1 parent 560f44b commit ddb2155

2 files changed

Lines changed: 128 additions & 0 deletions

File tree

‎datafusion/core/tests/sql/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ pub mod aggregates;
5353
pub mod create_drop;
5454
pub mod explain_analyze;
5555
pub mod joins;
56+
mod optimizer_convergence;
5657
mod path_partition;
5758
mod runtime_config;
5859
pub mod select;
Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3+
// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4+
5+
//! Checks that the logical optimizer reaches a fixed point: within `MAX`
6+
//! passes, a pass leaves the plan unchanged.
7+
//!
8+
//! The optimizer also stops when a pass produces a plan that an earlier pass
9+
//! produced (a cycle), so comparing the results of `max_passes = N` and
10+
//! `max_passes = N + 1` can not tell a fixed point from a cycle. The plan after
11+
//! each pass is recorded instead, and the last two must be equal.
12+
13+
use std::fs;
14+
use std::sync::Arc;
15+
16+
use datafusion::common::Result;
17+
use datafusion::datasource::MemTable;
18+
use datafusion::execution::SessionState;
19+
use datafusion::logical_expr::LogicalPlan;
20+
use datafusion::prelude::{SessionConfig, SessionContext};
21+
use test_utils::tpcds::tpcds_schemas;
22+
23+
const MAX: usize = 10;
24+
25+
fn ctx() -> SessionContext {
26+
let mut config = SessionConfig::default();
27+
config.options_mut().optimizer.max_passes = MAX;
28+
SessionContext::new_with_config(config)
29+
}
30+
31+
/// The plan before the first optimizer pass, followed by the plan after each
32+
/// pass that the optimizer ran.
33+
fn pass_outputs(state: &SessionState, plan: LogicalPlan) -> Result<Vec<String>> {
34+
let plan =
35+
state
36+
.analyzer()
37+
.execute_and_check(plan, state.config_options(), |_, _| {})?;
38+
let optimizer = state.optimizer();
39+
let last_rule = optimizer.rules.last().map(Arc::as_ptr);
40+
let mut out = vec![plan.display_indent().to_string()];
41+
optimizer.optimize(plan, state, |plan, rule| {
42+
if last_rule.is_some_and(|last| std::ptr::addr_eq(last, rule)) {
43+
out.push(plan.display_indent().to_string());
44+
}
45+
})?;
46+
Ok(out)
47+
}
48+
49+
/// The number of passes after which a pass left the plan unchanged, or `None`
50+
/// when the last pass still changed it (the plans still change at `MAX`, or
51+
/// the optimizer stopped at a cycle).
52+
fn passes_to_fixed_point(outputs: &[String]) -> Option<usize> {
53+
let [.., before, last] = outputs else {
54+
return None;
55+
};
56+
(before == last).then(|| outputs.len() - 2)
57+
}
58+
59+
async fn outputs_for(ctx: &SessionContext, sql: &str) -> Result<Vec<String>> {
60+
let (state, plan) = ctx.sql(sql).await?.into_parts();
61+
pass_outputs(&state, plan)
62+
}
63+
64+
async fn check(setup: &[&str], sql: &str) -> Result<()> {
65+
let ctx = ctx();
66+
for s in setup {
67+
ctx.sql(s).await?.collect().await?;
68+
}
69+
let outputs = outputs_for(&ctx, sql).await?;
70+
let n = passes_to_fixed_point(&outputs);
71+
assert!(
72+
n.is_some(),
73+
"did not converge:\n{}",
74+
outputs.last().unwrap()
75+
);
76+
Ok(())
77+
}
78+
79+
const SHAPE1_SETUP: &[&str] = &["CREATE TABLE test_data (id INT, float_val DOUBLE, \
80+
int_val INT, decimal_val DECIMAL(10,2)) AS VALUES (1, 5.3, 100, 100.00)"];
81+
const SHAPE1: &str = "SELECT * FROM test_data WHERE floor(int_val) = 100";
82+
83+
const SHAPE2_SETUP: &[&str] = &["CREATE TABLE t(a INT) AS VALUES (1),(2)"];
84+
const SHAPE2: &str = "SELECT a FROM t WHERE CASE WHEN arrow_field(a) IS NOT NULL \
85+
THEN arrow_field(a)['nullable'] IS NULL END";
86+
87+
#[tokio::test]
88+
async fn converges_floor_preimage() -> Result<()> {
89+
check(SHAPE1_SETUP, SHAPE1).await
90+
}
91+
92+
#[tokio::test]
93+
async fn converges_issue_25329() -> Result<()> {
94+
check(SHAPE2_SETUP, SHAPE2).await
95+
}
96+
97+
/// Checks that every TPC-DS query reaches a fixed point within `MAX` passes.
98+
#[tokio::test]
99+
async fn converges_tpcds() -> Result<()> {
100+
let ctx = ctx();
101+
for table in &tpcds_schemas() {
102+
ctx.register_table(
103+
table.name.as_str(),
104+
Arc::new(
105+
MemTable::try_new(Arc::new(table.schema.clone()), vec![vec![]])?
106+
.with_constraints(table.constraints.clone()),
107+
),
108+
)?;
109+
}
110+
let mut not_converged = vec![];
111+
for q in 1..=99 {
112+
let sql = fs::read_to_string(format!("tests/tpc-ds/{q}.sql"))
113+
.expect("Could not read query");
114+
// some queries have multiple statements
115+
for (i, sql) in sql.split(';').filter(|s| !s.trim().is_empty()).enumerate() {
116+
let outputs = outputs_for(&ctx, sql).await?;
117+
if passes_to_fixed_point(&outputs).is_none() {
118+
not_converged.push(format!("q{q}.{i}"));
119+
}
120+
}
121+
}
122+
assert!(
123+
not_converged.is_empty(),
124+
"not converged within {MAX} passes: {not_converged:?}"
125+
);
126+
Ok(())
127+
}

0 commit comments

Comments
 (0)