diff --git a/datafusion/core/tests/fuzz_cases/equivalence/ordering.rs b/datafusion/core/tests/fuzz_cases/equivalence/ordering.rs index 551d027999423..576d7e557dad3 100644 --- a/datafusion/core/tests/fuzz_cases/equivalence/ordering.rs +++ b/datafusion/core/tests/fuzz_cases/equivalence/ordering.rs @@ -3,10 +3,11 @@ // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements use crate::fuzz_cases::equivalence::utils::{ - NULL_PCTS, TestScalarUDF, assert_random_ordering_satisfy_is_sound, - contains_conservative_ordering_op, create_random_schema, create_test_params, - create_test_schema_2, generate_table_for_eq_properties, generate_table_for_orderings, - is_table_same_after_sort, + ALL_SORT_OPTIONS, NULL_PCTS, TestScalarUDF, assert_random_ordering_satisfy_is_sound, + contains_conservative_ordering_op, create_random_schema, + create_random_table_with_casts_and_windows, create_test_params, create_test_schema_2, + generate_table_for_eq_properties, generate_table_for_orderings, + is_table_same_after_sort, random_sort_options, }; use arrow::array::{ArrayRef, Float64Array, RecordBatch, new_null_array}; use arrow::compute::{SortOptions, sort}; @@ -23,6 +24,7 @@ use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; use itertools::Itertools; use rand::rngs::StdRng; +use rand::seq::SliceRandom; use rand::{Rng, SeedableRng}; use std::sync::Arc; @@ -235,24 +237,69 @@ fn test_ordering_satisfy_with_equivalence_complex_random() -> Result<()> { Ok(()) } -const ALL_SORT_OPTIONS: [SortOptions; 4] = [ - SortOptions { - descending: false, - nulls_first: false, - }, - SortOptions { - descending: false, - nulls_first: true, - }, - SortOptions { - descending: true, - nulls_first: false, - }, - SortOptions { - descending: true, - nulls_first: true, - }, -]; +/// Checks `ordering_satisfy` on tables with cast equalities and the orderings +/// of whole-partition windows (see `create_random_table_with_casts_and_windows`): +/// every ordering it claims must hold on the rows, and the orderings that +/// follow from the declared properties must be claimed. +#[test] +fn test_ordering_satisfy_with_casts_and_windows_random() -> Result<()> { + const N_RANDOM_TABLES: u64 = 100; + const N_ELEMENTS: usize = 125; + const N_DISTINCT: usize = 5; + const MAX_KEYS: usize = 4; + const SAMPLED_REQUIREMENT_PROBABILITY: f64 = 0.3; + + for (seed, &null_pct) in (0..N_RANDOM_TABLES).cartesian_product(NULL_PCTS) { + let table = create_random_table_with_casts_and_windows( + seed, null_pct, N_ELEMENTS, N_DISTINCT, + )?; + let eq_properties = &table.eq_properties; + let context = + format!("seed: {seed}, null_pct: {null_pct}, eq_properties: {eq_properties}"); + for ordering in &table.provable_orderings { + assert!( + is_table_same_after_sort(ordering.clone(), &table.batch)?, + "{context}, does not hold: {ordering}" + ); + assert!( + eq_properties.ordering_satisfy(ordering.clone())?, + "{context}, not proven: {ordering}" + ); + } + + let mut rng = StdRng::seed_from_u64(seed); + for n_keys in 1..=MAX_KEYS { + for mut indices in (0..table.keys.len()).combinations(n_keys) { + if n_keys > 2 && !rng.random_bool(SAMPLED_REQUIREMENT_PROBABILITY) { + continue; + } + indices.shuffle(&mut rng); + // Mostly in the direction each key is sorted in, so that many + // requirements hold and their claims get checked: + let sort_exprs = indices + .iter() + .map(|&idx| { + let (expr, options) = &table.keys[idx]; + let options = match options { + Some(options) if rng.random_bool(0.75) => *options, + _ => random_sort_options(&mut rng), + }; + PhysicalSortExpr::new(Arc::clone(expr), options) + }) + .collect::>(); + let ordering = LexOrdering::new(sort_exprs).unwrap(); + if eq_properties.ordering_satisfy(ordering.clone())? { + assert!( + is_table_same_after_sort(ordering.clone(), &table.batch)?, + "{context}, claimed: {ordering}" + ); + } + } + } + } + + Ok(()) +} /// Returns `0.0..=5.0` and `n_nulls` NULLs, sorted with `options`. fn sorted_f64_array(options: SortOptions, n_nulls: usize) -> Result { diff --git a/datafusion/core/tests/fuzz_cases/equivalence/properties.rs b/datafusion/core/tests/fuzz_cases/equivalence/properties.rs index 1d23335f39ddf..811bdc4ea597d 100644 --- a/datafusion/core/tests/fuzz_cases/equivalence/properties.rs +++ b/datafusion/core/tests/fuzz_cases/equivalence/properties.rs @@ -5,7 +5,8 @@ use std::sync::Arc; use crate::fuzz_cases::equivalence::utils::{ - NULL_PCTS, TestScalarUDF, create_random_schema, generate_table_for_eq_properties, + NULL_PCTS, TestScalarUDF, create_random_schema, + create_random_table_with_casts_and_windows, generate_table_for_eq_properties, is_table_same_after_sort, }; @@ -17,6 +18,9 @@ use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use datafusion_common::config::ConfigOptions; use itertools::Itertools; +use rand::rngs::StdRng; +use rand::seq::SliceRandom; +use rand::{Rng, SeedableRng}; #[test] fn test_find_longest_permutation_random() -> Result<()> { @@ -98,3 +102,62 @@ fn test_find_longest_permutation_random() -> Result<()> { Ok(()) } + +/// Checks `find_longest_permutation` on tables with cast equalities and the +/// orderings of whole-partition windows (see +/// `create_random_table_with_casts_and_windows`): the ordering it finds must +/// hold on the rows, and the sets of expressions that can be ordered in full +/// must be found in full. +#[test] +fn test_find_longest_permutation_with_casts_and_windows_random() -> Result<()> { + const N_RANDOM_TABLES: u64 = 100; + const N_ELEMENTS: usize = 125; + const N_DISTINCT: usize = 5; + const MAX_EXPRS: usize = 4; + const SAMPLED_SET_PROBABILITY: f64 = 0.3; + + for (seed, &null_pct) in (0..N_RANDOM_TABLES).cartesian_product(NULL_PCTS) { + let table = create_random_table_with_casts_and_windows( + seed, null_pct, N_ELEMENTS, N_DISTINCT, + )?; + let eq_properties = &table.eq_properties; + let context = + format!("seed: {seed}, null_pct: {null_pct}, eq_properties: {eq_properties}"); + let mut rng = StdRng::seed_from_u64(seed); + for mut exprs in table.provable_permutations.clone() { + exprs.shuffle(&mut rng); + let (_, indices) = eq_properties.find_longest_permutation(&exprs)?; + assert_eq!( + indices.len(), + exprs.len(), + "{context}, not found in full: [{}]", + exprs.iter().join(", ") + ); + } + + let candidates = table.keys.iter().map(|(expr, _)| expr).collect::>(); + for n_exprs in 1..=MAX_EXPRS { + for indices in (0..candidates.len()).combinations(n_exprs) { + if n_exprs > 2 && !rng.random_bool(SAMPLED_SET_PROBABILITY) { + continue; + } + let exprs = indices + .iter() + .map(|&idx| Arc::clone(candidates[idx])) + .collect::>(); + let (ordering, found) = eq_properties.find_longest_permutation(&exprs)?; + for (sort_expr, &idx) in ordering.iter().zip(&found) { + assert!(sort_expr.expr.eq(&exprs[idx]), "{context}"); + } + if let Some(ordering) = LexOrdering::new(ordering) { + assert!( + is_table_same_after_sort(ordering.clone(), &table.batch)?, + "{context}, found: {ordering}" + ); + } + } + } + } + + Ok(()) +} diff --git a/datafusion/core/tests/fuzz_cases/equivalence/utils.rs b/datafusion/core/tests/fuzz_cases/equivalence/utils.rs index 49da985a64935..4e536ce9be7bb 100644 --- a/datafusion/core/tests/fuzz_cases/equivalence/utils.rs +++ b/datafusion/core/tests/fuzz_cases/equivalence/utils.rs @@ -5,9 +5,12 @@ use std::cmp::Ordering; use std::sync::Arc; -use arrow::array::{ArrayRef, Float32Array, Float64Array, RecordBatch, UInt32Array}; +use arrow::array::{ + ArrayRef, AsArray, Float32Array, Float64Array, Int32Array, Int64Array, RecordBatch, + UInt32Array, +}; use arrow::compute::{SortColumn, SortOptions, lexsort_to_indices, take_record_batch}; -use arrow::datatypes::{DataType, Field, FieldRef, Schema, SchemaRef}; +use arrow::datatypes::{DataType, Field, FieldRef, Int32Type, Schema, SchemaRef}; use datafusion_common::tree_node::TreeNode; use datafusion_common::utils::{compare_rows, get_row_at_idx}; use datafusion_common::{Result, exec_err, internal_datafusion_err, plan_err}; @@ -22,7 +25,7 @@ use datafusion_physical_expr::equivalence::{ use datafusion_physical_expr::{ConstExpr, EquivalenceProperties}; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; -use datafusion_physical_plan::expressions::{BinaryExpr, Column, col}; +use datafusion_physical_plan::expressions::{BinaryExpr, CastExpr, Column, col}; use itertools::izip; use rand::prelude::*; @@ -548,8 +551,240 @@ pub fn generate_table_for_orderings( Ok(batch) } +/// A table, equivalence properties that hold on it, and sort keys to check +/// against both. See [`create_random_table_with_casts_and_windows`]. +pub struct TableWithProperties { + pub eq_properties: EquivalenceProperties, + pub batch: RecordBatch, + /// Candidate sort keys, each with the direction it is sorted in, if any. + pub keys: Vec<(Arc, Option)>, + /// Orderings that follow from the declared properties, which + /// `ordering_satisfy` should therefore prove. + pub provable_orderings: Vec, + /// Sets of expressions that `find_longest_permutation` should find in + /// full. + pub provable_permutations: Vec>>, +} + +/// Builds a random table, then declares on it only properties that hold, +/// including two kinds that [`create_random_schema`] does not generate: +/// equalities with casts, and the orderings that a window over a whole +/// partition adds. +/// +/// - `k`, `a` and `b` are sorted by `[k, a, b]` in random directions. +/// - `t` and `x` are `a` and `b` cast to `BIGINT`, and are declared equal to +/// these casts. `a` is a `DOUBLE`, so `t` keeps the order of `a` but not +/// its ties. The declared ordering is `[k, a, b]` or, at random, +/// `[k, t, a, b]`, from which `[k, a, b]` follows. The candidate keys also +/// cast `a` to `INT`, which loses its ties too. +/// - `w` has one value per `k`, like `SUM(..) OVER (PARTITION BY k)`, which +/// is declared as `[k, w]` in all four directions of `w`. +/// - `c` and `d` are sorted by `[c, d]`, independently of the rest. +/// - `e` is constant. +/// +/// The orderings are declared in random order. +pub fn create_random_table_with_casts_and_windows( + seed: u64, + null_pct: f64, + n_elem: usize, + n_distinct: usize, +) -> Result { + let mut rng = StdRng::seed_from_u64(seed); + let [opts_k, opts_a, opts_b, opts_c, opts_d] = + std::array::from_fn(|_| random_sort_options(&mut rng)); + + let k = generate_random_i32_array(n_elem, n_distinct, null_pct, &mut rng); + let a = generate_random_f64_array(n_elem, n_distinct, null_pct, &mut rng); + let b = generate_random_i32_array(n_elem, n_distinct, null_pct, &mut rng); + let [k, a, b] = lexsort_columns([(k, opts_k), (a, opts_a), (b, opts_b)])?; + let t = arrow::compute::cast(&a, &DataType::Int64)?; + let x = arrow::compute::cast(&b, &DataType::Int64)?; + // `k` is sorted, so rows with the same `k` are adjacent: + let mut w = Vec::with_capacity(n_elem); + let mut previous_k = None; + for value in k.as_primitive::().iter() { + if previous_k != Some(value) { + previous_k = Some(value); + w.push(random_value(n_distinct, null_pct, &mut rng).map(i64::from)); + } else { + w.push(*w.last().unwrap()); + } + } + let w = Arc::new(Int64Array::from(w)) as ArrayRef; + let c = generate_random_i32_array(n_elem, n_distinct, null_pct, &mut rng); + let d = generate_random_i32_array(n_elem, n_distinct, null_pct, &mut rng); + let [c, d] = lexsort_columns([(c, opts_c), (d, opts_d)])?; + let e = Arc::new(Int32Array::from(vec![1; n_elem])) as ArrayRef; + + let columns = [ + ("k", k), + ("a", a), + ("b", b), + ("t", t), + ("x", x), + ("w", w), + ("c", c), + ("d", d), + ("e", e), + ]; + let schema = Arc::new(Schema::new( + columns + .iter() + .map(|(name, array)| { + Field::new(*name, array.data_type().clone(), null_pct > 0.0) + }) + .collect::>(), + )); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + columns.into_iter().map(|(_, array)| array).collect(), + )?; + + let [k, a, b, t, x, w, c, d, e] = ["k", "a", "b", "t", "x", "w", "c", "d", "e"] + .map(|name| col(name, &schema).unwrap()); + let cast = |expr: &Arc, data_type| { + Arc::new(CastExpr::new(Arc::clone(expr), data_type, None)) + as Arc + }; + let cast_k = cast(&k, DataType::Int64); + let cast_a = cast(&a, DataType::Int64); + // In no equivalence class, unlike `CAST(a AS BIGINT)`: + let cast_a_int = cast(&a, DataType::Int32); + let cast_b = cast(&b, DataType::Int64); + let cast_w = cast(&w, DataType::Float64); + let sort = |expr: &Arc, options| { + PhysicalSortExpr::new(Arc::clone(expr), options) + }; + + let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + eq_properties.add_equal_conditions(Arc::clone(&t), Arc::clone(&cast_a))?; + eq_properties.add_equal_conditions(Arc::clone(&x), Arc::clone(&cast_b))?; + eq_properties.add_constants([ConstExpr::from(Arc::clone(&e))])?; + let mut main_ordering = vec![sort(&k, opts_k), sort(&a, opts_a), sort(&b, opts_b)]; + if rng.random_bool(0.5) { + main_ordering.insert(1, sort(&t, opts_a)); + } + let mut orderings = vec![main_ordering, vec![sort(&c, opts_c), sort(&d, opts_d)]]; + orderings.extend( + ALL_SORT_OPTIONS.map(|options| vec![sort(&k, opts_k), sort(&w, options)]), + ); + orderings.shuffle(&mut rng); + eq_properties.add_orderings(orderings); + + let keys = vec![ + (Arc::clone(&k), Some(opts_k)), + (Arc::clone(&a), Some(opts_a)), + (Arc::clone(&b), Some(opts_b)), + (Arc::clone(&t), Some(opts_a)), + (Arc::clone(&x), Some(opts_b)), + (Arc::clone(&w), None), + (Arc::clone(&c), Some(opts_c)), + (Arc::clone(&d), Some(opts_d)), + (Arc::clone(&e), None), + (Arc::clone(&cast_k), Some(opts_k)), + (Arc::clone(&cast_a), Some(opts_a)), + (Arc::clone(&cast_a_int), Some(opts_a)), + (Arc::clone(&cast_b), Some(opts_b)), + (Arc::clone(&cast_w), None), + ]; + // `x` and `CAST(b AS BIGINT)` both normalize to the member that + // represents their class, which no ordering leads, so neither is proven + // to be ordered after `[k, a]`. The same holds for `t` and + // `CAST(a AS BIGINT)` when the declared ordering is `[k, a, b]`. They are + // still checked as candidate keys. + let provable_orderings = [ + vec![sort(&k, opts_k), sort(&a, opts_a), sort(&b, opts_b)], + vec![sort(&cast_k, opts_k), sort(&a, opts_a)], + vec![sort(&e, opts_c), sort(&c, opts_c), sort(&d, opts_d)], + ] + .into_iter() + .chain(ALL_SORT_OPTIONS.into_iter().flat_map(|options| { + [ + vec![ + sort(&k, opts_k), + sort(&w, options), + sort(&a, opts_a), + sort(&b, opts_b), + ], + vec![sort(&k, opts_k), sort(&cast_w, options), sort(&a, opts_a)], + ] + })) + .map(|ordering| LexOrdering::new(ordering).unwrap()) + .collect::>(); + let provable_permutations = vec![ + vec![ + Arc::clone(&w), + Arc::clone(&b), + Arc::clone(&a), + Arc::clone(&k), + ], + vec![cast_w, b, a, k], + ]; + + Ok(TableWithProperties { + eq_properties, + batch, + keys, + provable_orderings, + provable_permutations, + }) +} + +/// Sorts `columns` together, the first one first, each in its direction. +fn lexsort_columns( + columns: [(ArrayRef, SortOptions); N], +) -> Result<[ArrayRef; N]> { + let sort_columns = columns.map(|(values, options)| SortColumn { + values, + options: Some(options), + }); + let sorted = arrow::compute::lexsort(&sort_columns, None)?; + Ok(sorted.try_into().unwrap()) +} + +/// The four directions of a sort key. +pub const ALL_SORT_OPTIONS: [SortOptions; 4] = [ + SortOptions { + descending: false, + nulls_first: false, + }, + SortOptions { + descending: false, + nulls_first: true, + }, + SortOptions { + descending: true, + nulls_first: false, + }, + SortOptions { + descending: true, + nulls_first: true, + }, +]; + pub const NULL_PCTS: &[f64] = &[0.0, 0.1, 0.5]; +/// A value in `0..n_distinct`, or `None` with probability `null_pct`. +fn random_value(n_distinct: usize, null_pct: f64, rng: &mut StdRng) -> Option { + if rng.random::() < null_pct { + None + } else { + Some(rng.random_range(0..n_distinct as i32)) + } +} + +fn generate_random_i32_array( + n_elems: usize, + n_distinct: usize, + null_pct: f64, + rng: &mut StdRng, +) -> ArrayRef { + let values = (0..n_elems) + .map(|_| random_value(n_distinct, null_pct, rng)) + .collect::(); + Arc::new(values) +} + // Utility function to generate random f64 array fn generate_random_f64_array( n_elems: usize, diff --git a/datafusion/physical-expr/benches/equivalence_properties.rs b/datafusion/physical-expr/benches/equivalence_properties.rs index 5846b91ed8b62..b9318871d2235 100644 --- a/datafusion/physical-expr/benches/equivalence_properties.rs +++ b/datafusion/physical-expr/benches/equivalence_properties.rs @@ -26,7 +26,8 @@ use std::sync::Arc; use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use criterion::{BatchSize, BenchmarkId, Criterion, criterion_group, criterion_main}; -use datafusion_physical_expr::expressions::Column; +use datafusion_expr::Operator; +use datafusion_physical_expr::expressions::{BinaryExpr, Column, lit}; use datafusion_physical_expr::{ EquivalenceProperties, LexOrdering, PhysicalExpr, PhysicalSortExpr, PhysicalSortRequirement, @@ -84,6 +85,21 @@ fn bench_ordering_satisfaction(c: &mut Criterion) { for n_classes in [2, 8, 32] { let props = properties(n_classes); + // Check each answer once, untimed, so that a wrong answer cannot pass + // for a faster one. + assert!(props.ordering_satisfy(one_key.clone()).unwrap()); + assert!(!props.ordering_satisfy(one_key_unsatisfied.clone()).unwrap()); + assert!(props.ordering_satisfy(four_keys.clone()).unwrap()); + assert!(props.ordering_satisfy_requirement(one_req.clone()).unwrap()); + assert!( + props + .ordering_satisfy_requirement(four_reqs.clone()) + .unwrap() + ); + let (_, all_satisfied) = props + .extract_common_sort_prefix(four_key_ordering.clone()) + .unwrap(); + assert!(all_satisfied); // A single sort key: the most common shape by far. group.bench_function( @@ -106,7 +122,8 @@ fn bench_ordering_satisfaction(c: &mut Criterion) { ) }, ); - // Four sort keys: exercises the per-key constant registration. + // Four sort keys: each key only has to be sorted among the rows that tie + // on the keys before it. group.bench_function( BenchmarkId::new("ordering_satisfy/4_keys", n_classes), |b| { @@ -152,5 +169,94 @@ fn bench_ordering_satisfaction(c: &mut Criterion) { group.finish(); } -criterion_group!(benches, bench_ordering_satisfaction); +/// Other shapes: `find_longest_permutation`, the orderings that window +/// aggregates add, and a deeply nested sort key. +fn bench_shapes(c: &mut Criterion) { + let mut group = c.benchmark_group("equivalence_properties_shapes"); + let opts = [(false, false), (false, true), (true, false), (true, true)].map( + |(descending, nulls_first)| SortOptions { + descending, + nulls_first, + }, + ); + let asc_nulls_last = opts[0]; + + // `find_longest_permutation` with one expression, and with five: `c0`, + // `c1`, `c2` and `c3` are found one per round, `c7` never. + for n_classes in [2, 32] { + let props = properties(n_classes); + let one = vec![col(0)]; + let five = vec![col(3), col(2), col(1), col(0), col(7)]; + assert_eq!(props.find_longest_permutation(&one).unwrap().1, [0]); + assert_eq!( + props.find_longest_permutation(&five).unwrap().1, + [3, 2, 1, 0] + ); + group.bench_function( + BenchmarkId::new("find_longest_permutation/1_expr", n_classes), + |b| b.iter(|| props.find_longest_permutation(&one).unwrap()), + ); + group.bench_function( + BenchmarkId::new("find_longest_permutation/5_exprs", n_classes), + |b| b.iter(|| props.find_longest_permutation(&five).unwrap()), + ); + } + + // `w` window aggregates over a whole partition, such as + // `SUM(x) OVER (PARTITION BY c0)`: each is constant per `c0`, which is + // recorded by adding `[c0, w_i]` in all four directions of `w_i` + // (columns `c1..=cw`). The requirement is `[c0, c1]`. + for w in [4, 16, 32] { + let mut props = EquivalenceProperties::new(schema(w + 1)); + let mut orderings = vec![]; + for i in 1..=w { + for o in opts { + orderings.push(vec![ + PhysicalSortExpr::new(col(0), asc_nulls_last), + PhysicalSortExpr::new(col(i), o), + ]); + } + } + props.add_orderings(orderings); + let keys = vec![ + PhysicalSortExpr::new(col(0), asc_nulls_last), + PhysicalSortExpr::new(col(1), asc_nulls_last), + ]; + assert!(props.ordering_satisfy(keys.clone()).unwrap()); + group.bench_function(BenchmarkId::new("window", w), |b| { + b.iter_batched( + || keys.clone(), + |keys| props.ordering_satisfy(keys).unwrap(), + BatchSize::SmallInput, + ) + }); + } + + // Deep key: `[c0, e, c1]` with `e = c1 + 1 + ... + 1` nested `depth` deep. + // `+` may overflow, so `e` is unordered: the check computes the properties + // of `e` and stops there. + for depth in [4, 8, 16, 32] { + let props = properties(8); + let mut e = col(1); + for _ in 0..depth { + e = Arc::new(BinaryExpr::new(e, Operator::Plus, lit(1i32))); + } + let keys = vec![ + asc(0), + PhysicalSortExpr::new(e, SortOptions::default()), + asc(1), + ]; + assert!(!props.ordering_satisfy(keys.clone()).unwrap()); + group.bench_function(BenchmarkId::new("deep_key", depth), |b| { + b.iter_batched( + || keys.clone(), + |keys| props.ordering_satisfy(keys).unwrap(), + BatchSize::SmallInput, + ) + }); + } + group.finish(); +} + +criterion_group!(benches, bench_ordering_satisfaction, bench_shapes); criterion_main!(benches); diff --git a/datafusion/physical-expr/src/equivalence/ordering.rs b/datafusion/physical-expr/src/equivalence/ordering.rs index f0542d9a399a1..7de47767f62fd 100644 --- a/datafusion/physical-expr/src/equivalence/ordering.rs +++ b/datafusion/physical-expr/src/equivalence/ordering.rs @@ -12,7 +12,7 @@ use crate::{LexOrdering, PhysicalExpr, add_offset_to_physical_sort_exprs}; use arrow::compute::SortOptions; use arrow::datatypes::SchemaRef; -use datafusion_common::{HashSet, Result}; +use datafusion_common::Result; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; /// An `OrderingEquivalenceClass` keeps track of distinct alternative orderings @@ -210,25 +210,16 @@ impl OrderingEquivalenceClass { None } - /// Checks whether the given expression is partially constant according to - /// this ordering equivalence class. + /// Checks whether `expr` has a single value according to this ordering + /// equivalence class: whether one ordering starts with `expr` in some + /// direction and another starts with it in the exact reverse, such as + /// `ASC NULLS LAST` and `DESC NULLS FIRST`. Rows sorted both ways all have + /// the same value (or are all null). /// - /// This function determines whether `expr` appears in at least one combination - /// of `descending` and `nulls_first` options that indicate partial constantness - /// in a lexicographical ordering. Specifically, an expression is considered - /// a partial constant in this context if its `SortOptions` satisfies either - /// of the following conditions: - /// - It is `descending` with `nulls_first` and _also_ `ascending` with - /// `nulls_last`, OR - /// - It is `descending` with `nulls_last` and _also_ `ascending` with - /// `nulls_first`. - /// - /// The equivalence mechanism primarily uses `ConstExpr`s to represent globally - /// constant expressions. However, some expressions may only be partially - /// constant within a lexicographical ordering. This function helps identify - /// such cases. If an expression is constant within a prefix ordering, it is - /// added as a constant during `ordering_satisfy_requirement()` iterations - /// after the corresponding prefix requirement is satisfied. + /// This is how a value that is constant only within groups of rows is + /// recorded, since a `ConstExpr` is constant across all rows. For example, + /// `SUM(x) OVER (PARTITION BY k)` has one value per `k`, so the window adds + /// `[k, w]` once for each of the four directions of `w`. /// /// ### Future Improvements /// @@ -238,29 +229,30 @@ impl OrderingEquivalenceClass { /// 2. `EquivalenceProperties` gains `FunctionalDependency` awareness, eliminating /// the need for `Constant` and `Constraints`. pub fn is_expr_partial_const(&self, expr: &Arc) -> bool { - let mut constantness_defining_pairs = [ - HashSet::from([(false, false), (true, true)]), - HashSet::from([(false, true), (true, false)]), - ]; - - for ordering in self.iter() { - let leading_ordering = ordering.first(); - if leading_ordering.expr.eq(expr) { - let opt = ( - leading_ordering.options.descending, - leading_ordering.options.nulls_first, - ); - constantness_defining_pairs[0].remove(&opt); - constantness_defining_pairs[1].remove(&opt); - } - } - - constantness_defining_pairs - .iter() - .any(|pair| pair.is_empty()) + is_partial_const(self.iter().map(|ordering| ordering.first()), expr) } } +/// Checks whether `expr` has a single value according to the leading sort +/// expressions of some orderings: whether `expr` leads one of them in some +/// direction and another in the exact reverse. See +/// [`OrderingEquivalenceClass::is_expr_partial_const`]. +pub(crate) fn is_partial_const<'a>( + leading: impl Iterator + Clone, + expr: &Arc, +) -> bool { + let leads_in = move |options: SortOptions| { + leading + .clone() + .any(|sort_expr| sort_expr.options == options && sort_expr.expr.eq(expr)) + }; + // `!` gives the exact reverse, e.g. `DESC NULLS FIRST` for `ASC NULLS LAST`. + let asc_nulls_first = SortOptions::default().asc().nulls_first(); + let asc_nulls_last = SortOptions::default().asc().nulls_last(); + (leads_in(asc_nulls_first) && leads_in(!asc_nulls_first)) + || (leads_in(asc_nulls_last) && leads_in(!asc_nulls_last)) +} + impl Deref for OrderingEquivalenceClass { type Target = [LexOrdering]; @@ -978,4 +970,64 @@ mod tests { Ok(()) } + + #[test] + fn test_is_expr_partial_const() -> Result<()> { + let schema = create_test_schema()?; + let col_a = &col("a", &schema)?; + let col_b = &col("b", &schema)?; + let [ + asc_nulls_last, + asc_nulls_first, + desc_nulls_last, + desc_nulls_first, + ] = [(false, false), (false, true), (true, false), (true, true)].map( + |(descending, nulls_first)| SortOptions { + descending, + nulls_first, + }, + ); + + // The leading sort expression of each ordering, and whether `a` has a + // single value according to them. + let test_cases = vec![ + // A direction and its exact reverse: + ( + vec![(col_a, asc_nulls_last), (col_a, desc_nulls_first)], + true, + ), + ( + vec![(col_a, asc_nulls_first), (col_a, desc_nulls_last)], + true, + ), + // Nulls first both ways, e.g. `NULL, NULL, 5, 5`: + ( + vec![(col_a, asc_nulls_first), (col_a, desc_nulls_first)], + false, + ), + // Nulls last both ways, e.g. `5, 5, NULL, NULL`: + ( + vec![(col_a, asc_nulls_last), (col_a, desc_nulls_last)], + false, + ), + (vec![(col_a, asc_nulls_last)], false), + // The reverse direction belongs to another expression: + ( + vec![(col_a, asc_nulls_last), (col_b, desc_nulls_first)], + false, + ), + ]; + for (leading, expected) in test_cases { + let oeq_class = OrderingEquivalenceClass::new(leading.into_iter().map( + |(expr, options)| [PhysicalSortExpr::new(Arc::clone(expr), options)], + )); + assert_eq!( + oeq_class.is_expr_partial_const(col_a), + expected, + "{oeq_class}" + ); + } + + Ok(()) + } } diff --git a/datafusion/physical-expr/src/equivalence/properties/mod.rs b/datafusion/physical-expr/src/equivalence/properties/mod.rs index 4d21480a817f1..d6ba1ef18e8bb 100644 --- a/datafusion/physical-expr/src/equivalence/properties/mod.rs +++ b/datafusion/physical-expr/src/equivalence/properties/mod.rs @@ -5,6 +5,7 @@ mod dependency; // Submodule containing DependencyMap and Dependencies mod joins; // Submodule containing join_equivalence_properties mod union; // Submodule containing calculate_union +mod walk; // Submodule containing SortKeyWalk pub use joins::*; pub use union::*; @@ -17,6 +18,7 @@ use self::dependency::{ Dependencies, DependencyMap, construct_prefix_orderings, generate_dependency_orderings, referred_dependencies, }; +use self::walk::SortKeyWalk; use crate::equivalence::{ AcrossPartitions, EquivalenceGroup, OrderingEquivalenceClass, ProjectionMapping, }; @@ -28,15 +30,12 @@ use crate::{ use arrow::compute::SortOptions; use arrow::datatypes::SchemaRef; -use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion_common::{Constraint, Constraints, HashMap, Result, plan_err}; use datafusion_expr::interval_arithmetic::Interval; use datafusion_expr::sort_properties::{ExprProperties, SortProperties}; -use datafusion_physical_expr_common::sort_expr::options_compatible; -use datafusion_physical_expr_common::utils::ExprPropertiesNode; use indexmap::IndexSet; -use itertools::Itertools; +use itertools::{Either, Itertools}; /// `EquivalenceProperties` stores information about the output of a plan node /// that can be used to optimize the plan. Currently, it keeps track of: @@ -390,7 +389,7 @@ impl EquivalenceProperties { // Discover any new orderings based on the new equivalence classes: let leading_exprs: Vec<_> = self.oeq_cache.leading_map.keys().cloned().collect(); for expr in leading_exprs { - self.discover_new_orderings(expr)?; + self.discover_new_orderings(&expr)?; } Ok(()) } @@ -458,83 +457,96 @@ impl EquivalenceProperties { // Discover any new orderings based on the constants: let leading_exprs: Vec<_> = self.oeq_cache.leading_map.keys().cloned().collect(); for expr in leading_exprs { - self.discover_new_orderings(expr)?; + self.discover_new_orderings(&expr)?; } Ok(()) } /// Discover new valid orderings in light of a new equality. Accepts a single - /// argument (`expr`) which is used to determine the orderings to update. + /// argument (`normal_expr`) which is used to determine the orderings to update. /// When constants or equivalence classes change, there may be new orderings /// that can be discovered with the new equivalence properties. /// For a discussion, see: fn discover_new_orderings( &mut self, - normal_expr: Arc, + normal_expr: &Arc, ) -> Result<()> { - let Some(ordering_idxs) = self.oeq_cache.leading_map.get(&normal_expr) else { + let Some(ordering_idxs) = self.oeq_cache.leading_map.get(normal_expr) else { return Ok(()); }; - let eq_class = self - .eq_group - .get_equivalence_class(&normal_expr) - .map_or_else(|| vec![normal_expr], |class| class.clone().into()); - let mut new_orderings = vec![]; for idx in ordering_idxs { let ordering = &self.oeq_cache.normal_cls[*idx]; - let leading_ordering_options = ordering[0].options; - - 'exprs: for equivalent_expr in &eq_class { - let children = equivalent_expr.children(); - if children.is_empty() { - continue; - } - // Check if all children match the next expressions in the ordering: - let mut child_properties = vec![]; - // Build properties for each child based on the next expression: - for (i, child) in children.into_iter().enumerate() { - let Some(next) = ordering.get(i + 1) else { - break 'exprs; - }; - if !next.expr.eq(child) { - break 'exprs; - } - let data_type = child.data_type(&self.schema)?; - child_properties.push(ExprProperties { - sort_properties: SortProperties::Ordered(next.options), - range: Interval::make_unbounded(&data_type)?, - preserves_lex_ordering: true, - strictly_order_preserving: true, - }); - } - // Check if the expression is monotonic in all arguments. If - // that cannot be derived, no new ordering is found: - let Ok(expr_properties) = - equivalent_expr.get_properties(&child_properties) - else { - continue; - }; - if expr_properties.preserves_lex_ordering - && expr_properties.sort_properties - == SortProperties::Ordered(leading_ordering_options) - { - // Assume that `[c ASC, a ASC, b ASC]` is among existing - // orderings. If equality `c = f(a, b)` is given, ordering - // `[a ASC, b ASC]` implies the ordering `[c ASC]`. Thus, - // ordering `[a ASC, b ASC]` is also a valid ordering. - new_orderings.push(ordering[1..].to_vec()); - break; - } - } + new_orderings + .extend(self.implied_ordering(&ordering[0], ordering[1..].iter())?); } - if !new_orderings.is_empty() { self.add_orderings(new_orderings); } Ok(()) } + /// Returns the ordering that the normal ordering `[leading, rest..]` + /// implies through an equality on its leading expression, if any. + /// + /// If data is sorted by `[c, a, b]` and `c = f(a, b)` for a function `f` + /// that preserves the lexicographical order of its arguments, the data is + /// also sorted by `[a, b]`: rows with different `(a, b)` are in the order + /// of `f(a, b)`, which is the order of `c`, and rows with equal `c` are + /// sorted by `[a, b]` anyway. Whatever follows `b` stays sorted as well, so + /// the implied ordering is all of `rest`. + fn implied_ordering<'b>( + &self, + leading: &PhysicalSortExpr, + rest: impl Iterator + Clone, + ) -> Result>> { + // The expressions equal to `leading`; any of them can play `f(a, b)`: + let members = match self.eq_group.get_equivalence_class(&leading.expr) { + Some(class) => Either::Left(class.iter()), + None => Either::Right(std::iter::once(&leading.expr)), + }; + for member in members { + let args = member.children(); + // An expression without arguments, such as a column, implies + // nothing: + if args.is_empty() { + continue; + } + // The arguments must be the next expressions of the ordering, in + // order, and are then sorted the way the ordering says: + let mut sorted = rest.clone(); + let mut arg_properties = vec![]; + for arg in args { + match sorted.next() { + Some(sort_expr) if sort_expr.expr.eq(arg) => { + let data_type = arg.data_type(&self.schema)?; + arg_properties.push(ExprProperties { + sort_properties: SortProperties::Ordered(sort_expr.options), + range: Interval::make_unbounded(&data_type)?, + preserves_lex_ordering: true, + strictly_order_preserving: true, + }); + } + // This ends the whole search: members after this one are + // not tried, even if their arguments would match. + _ => return Ok(None), + } + } + // `member` must preserve the lexicographical order of its + // arguments, in the direction of `leading`. If that cannot be + // derived, this member implies nothing: + let Ok(properties) = member.get_properties(&arg_properties) else { + continue; + }; + if properties.preserves_lex_ordering + && properties.sort_properties == SortProperties::Ordered(leading.options) + { + return Ok(Some(rest.cloned().collect())); + } + } + Ok(None) + } + /// Updates the ordering equivalence class within assuming that the table /// is re-sorted according to the argument `ordering`, and returns whether /// this operation resulted in any change. Note that equivalence classes @@ -609,8 +621,9 @@ impl EquivalenceProperties { /// **Existing orderings:** `[[a, b, c], [a, d]]`, **constants:** `[]` /// 1. The function first checks the leading requirement `a`, which is /// satisfied by `[a, b, c].first()`. - /// 2. `a` is added as a constant for the next iteration. - /// 3. Normal orderings become `[[b, c], [d]]`. + /// 2. `a` is pinned: the next key only has to be sorted among rows that + /// tie on `a`. + /// 3. The orderings now start at `[b, c]` and `[d]`. /// 4. The function fails for `c` in the second iteration, as neither /// `[b, c]` nor `[d]` satisfies `c`. /// @@ -619,8 +632,9 @@ impl EquivalenceProperties { /// **Existing orderings:** `[[a, b, c], [a, d]]`, **constants:** `[]` /// 1. The function first checks the leading requirement `a`, which is /// satisfied by `[a, b, c].first()`. - /// 2. `a` is added as a constant for the next iteration. - /// 3. Normal orderings become `[[b, c], [d]]`. + /// 2. `a` is pinned: the next key only has to be sorted among rows that + /// tie on `a`. + /// 3. The orderings now start at `[b, c]` and `[d]`. /// 4. The function returns `true` as `[d]` satisfies `d`. pub fn ordering_satisfy_requirement( &self, @@ -635,83 +649,8 @@ impl EquivalenceProperties { if self.satisfied_by_constraints(&normal_reqs) { return Ok(true); } - let schema = self.schema(); - // Registering satisfied keys as constants mutates the state, so it - // needs an owned copy -- but only from the second requirement onwards. - // Single-element requirements (the common case) never pay for the clone. - let last_idx = normal_reqs.len() - 1; - let mut owned = None::; - for (idx, element) in normal_reqs.into_iter().enumerate() { - let eq_properties = owned.as_ref().unwrap_or(self); - // Check whether given requirement is satisfied: - let ExprProperties { - sort_properties, .. - } = eq_properties.get_expr_properties(Arc::clone(&element.expr)); - let satisfy = match sort_properties { - SortProperties::Ordered(options) => element.options.is_none_or(|opts| { - let nullable = element.expr.nullable(schema).unwrap_or(true); - options_compatible(&options, &opts, nullable) - }), - // Singleton expressions satisfy any requirement. - SortProperties::Singleton => true, - SortProperties::Unordered => false, - }; - if !satisfy { - return Ok(false); - } - if idx == last_idx { - // Nothing left to check, so no need to update the state: - break; - } - // Treat satisfied keys (and the sub-expressions they pin down) as - // constants in subsequent iterations. See - // [`Self::add_satisfied_key_constants`] for the rationale. - owned - .get_or_insert_with(|| self.clone()) - .add_satisfied_key_constants(element.expr)?; - } - Ok(true) - } - - /// Registers a satisfied sort key as a constant for subsequent iterations - /// of the ordering satisfaction checks. We can do this because the "next" - /// key only matters in a lexicographical ordering when the keys to its - /// left have the same values (i.e. within a single tie group). Note that - /// these expressions are not properly "constants"; this is just an - /// implementation strategy confined to the satisfaction checks. - /// - /// For example, assume that the requirement is `[a ASC, (b + c) ASC]`, - /// and existing equivalent orderings are `[a ASC, b ASC]` and `[c ASC]`. - /// Once we deduce that `[a ASC]` is satisfied, we add column `a` as a - /// constant to the algorithm state. This enables us to deduce that - /// `(b + c) ASC` is satisfied, given `a` is constant. - /// - /// In addition to the key itself, this also registers any sub-expressions - /// whose values the key pins down: if an expression is strictly - /// order-preserving, equal outputs imply equal values of its ordered - /// children, so within a tie group of the key those children are constant - /// as well. For example, if data is sorted by `[a, b]`, the requirement - /// `[CAST(a AS BIGINT) ASC, b ASC]` is satisfied: `a` is constant within - /// each group of equal `CAST(a AS BIGINT)` values, and hence `b` is - /// sorted within each such group. - fn add_satisfied_key_constants(&mut self, expr: Arc) -> Result<()> { - let mut stack = vec![expr]; - while let Some(expr) = stack.pop() { - let properties = self.get_expr_properties(Arc::clone(&expr)); - if properties.strictly_order_preserving { - for child in expr.children() { - let child_properties = self.get_expr_properties(Arc::clone(child)); - if matches!( - child_properties.sort_properties, - SortProperties::Ordered(_) - ) { - stack.push(Arc::clone(child)); - } - } - } - self.add_constants(std::iter::once(ConstExpr::from(expr)))?; - } - Ok(()) + let keys = normal_reqs.iter().map(|req| (&req.expr, req.options)); + Ok(SortKeyWalk::new(self).sorted_prefix_len(keys)? == normal_reqs.len()) } /// Returns the number of consecutive sort expressions (starting from the @@ -724,45 +663,10 @@ impl EquivalenceProperties { // length: return Ok(full_length); } - let schema = self.schema(); - // Registering satisfied keys as constants mutates the state, so it - // needs an owned copy -- but only from the second sort expression - // onwards. Single-element orderings never pay for the clone. - let mut owned = None::; - for (idx, element) in normal_ordering.into_iter().enumerate() { - let eq_properties = owned.as_ref().unwrap_or(self); - // Check whether given ordering is satisfied: - let ExprProperties { - sort_properties, .. - } = eq_properties.get_expr_properties(Arc::clone(&element.expr)); - let satisfy = match sort_properties { - SortProperties::Ordered(options) => options_compatible( - &options, - &element.options, - element.expr.nullable(schema).unwrap_or(true), - ), - // Singleton expressions satisfy any ordering. - SortProperties::Singleton => true, - SortProperties::Unordered => false, - }; - if !satisfy { - // As soon as one sort expression is unsatisfied, return how - // many we've satisfied so far: - return Ok(idx); - } - if idx + 1 == full_length { - // Nothing left to check, so no need to update the state: - break; - } - // Treat satisfied keys (and the sub-expressions they pin down) as - // constants in subsequent iterations. See - // [`Self::add_satisfied_key_constants`] for the rationale. - owned - .get_or_insert_with(|| self.clone()) - .add_satisfied_key_constants(Arc::clone(&element.expr))?; - } - // All sort expressions are satisfied, return full length: - Ok(full_length) + let keys = normal_ordering + .iter() + .map(|sort_expr| (&sort_expr.expr, Some(sort_expr.options))); + SortKeyWalk::new(self).sorted_prefix_len(keys) } /// Determines the longest normal prefix of `ordering` satisfied by the @@ -1301,13 +1205,12 @@ impl EquivalenceProperties { &self, exprs: &[Arc], ) -> Result<(Vec, Vec)> { - let mut eq_properties = self.clone(); + let mut walk = SortKeyWalk::new(self); let mut result = vec![]; // The algorithm is as follows: // - Iterate over all the expressions and insert ordered expressions // into the result. - // - Treat inserted expressions as constants (i.e. add them as constants - // to the state). + // - Pin inserted expressions (see [`SortKeyWalk`]). // - Continue the above procedure until no expression is inserted; i.e. // the algorithm reaches a fixed point. // This algorithm should reach a fixed point in at most `exprs.len()` @@ -1320,7 +1223,7 @@ impl EquivalenceProperties { .filter_map(|&idx| { let ExprProperties { sort_properties, .. - } = eq_properties.get_expr_properties(Arc::clone(&exprs[idx])); + } = walk.expr_properties(Arc::clone(&exprs[idx])); match sort_properties { SortProperties::Ordered(options) => { let expr = Arc::clone(&exprs[idx]); @@ -1340,17 +1243,18 @@ impl EquivalenceProperties { break; } // Remove indices that have an ordering from `search_indices`, and - // treat ordered expressions as constants in subsequent iterations. - // We can do this because the "next" key only matters in a lexicographical + // pin the ordered expressions (in normal form, as the walk sees + // the orderings) for the next iterations, if any. We can do this + // because the "next" key only matters in a lexicographical // ordering when the keys to its left have the same values. - // - // Note that these expressions are not properly "constants". This is just - // an implementation strategy confined to this function. - for (PhysicalSortExpr { expr, .. }, idx) in &ordered_exprs { - let const_expr = ConstExpr::from(Arc::clone(expr)); - eq_properties.add_constants(std::iter::once(const_expr))?; + for (_, idx) in &ordered_exprs { search_indices.shift_remove(idx); } + if !search_indices.is_empty() { + for (PhysicalSortExpr { expr, .. }, _) in &ordered_exprs { + walk.pin(self.eq_group.normalize_expr(Arc::clone(expr)))?; + } + } // Add new ordered section to the state. result.extend(ordered_exprs); } @@ -1396,11 +1300,7 @@ impl EquivalenceProperties { /// Returns an [`ExprProperties`] object containing the ordering and range /// information for the given expression. pub fn get_expr_properties(&self, expr: Arc) -> ExprProperties { - ExprPropertiesNode::new_unknown(expr) - .transform_up(|expr| update_properties(expr, self)) - .data() - .map(|node| node.data) - .unwrap_or_else(|_| ExprProperties::new_unknown()) + SortKeyWalk::new(self).expr_properties(expr) } /// Returns whether `transformed_expr` reports the same ordering direction @@ -1524,68 +1424,6 @@ impl Display for EquivalenceProperties { } } -/// Calculates the properties of a given [`ExprPropertiesNode`]. -/// -/// Order information can be retrieved as: -/// - If it is a leaf node, we directly find the order of the node by looking -/// at the given sort expression and equivalence properties if it is a `Column` -/// leaf, or we mark it as unordered. In the case of a `Literal` leaf, we mark -/// it as singleton so that it can cooperate with all ordered columns. -/// - If it is an intermediate node, the children states matter. Each `PhysicalExpr` -/// and operator has its own rules on how to propagate the children orderings. -/// However, before we engage in recursion, we check whether this intermediate -/// node directly matches with the sort expression. If there is a match, the -/// sort expression emerges at that node immediately, discarding the recursive -/// result coming from its children. -/// -/// Range information is calculated as: -/// - If it is a `Literal` node, we set the range as a point value. If it is a -/// `Column` node, we set the datatype of the range, but cannot give an interval -/// for the range, yet. -/// - If it is an intermediate node, the children states matter. Each `PhysicalExpr` -/// and operator has its own rules on how to propagate the children range. -fn update_properties( - mut node: ExprPropertiesNode, - eq_properties: &EquivalenceProperties, -) -> Result> { - // First, try to gather the information from the children: - if !node.expr.children().is_empty() { - // We have an intermediate (non-leaf) node, account for its children: - let children_props = node.children.iter().map(|c| c.data.clone()).collect_vec(); - // Properties that cannot be derived from the children, for example of a - // function whose input may be outside its domain, are unknown. The - // ordering of the node itself can still be known, so do not fail here, - // and keep the type of its range so that parents can derive from it. - node.data = node.expr.get_properties(&children_props).or_else(|_| { - Interval::make_unbounded(&node.expr.data_type(eq_properties.schema())?) - .map(|range| ExprProperties::new_unknown().with_range(range)) - })?; - } else if node.expr.is::() { - // We have a Literal, which is one of the two possible leaf node types: - node.data = node.expr.get_properties(&[])?; - } else if node.expr.is::() { - // We have a Column, which is the other possible leaf node type: - node.data.range = - Interval::make_unbounded(&node.expr.data_type(eq_properties.schema())?)?; - // A column is the identity mapping of itself, which is trivially - // strict: - node.data.strictly_order_preserving = true; - } - // Now, check what we know about orderings: - let normal_expr = eq_properties - .eq_group - .normalize_expr(Arc::clone(&node.expr)); - let oeq_class = &eq_properties.oeq_cache.normal_cls; - if eq_properties.is_expr_constant(&normal_expr).is_some() - || oeq_class.is_expr_partial_const(&normal_expr) - { - node.data.sort_properties = SortProperties::Singleton; - } else if let Some(options) = oeq_class.get_options(&normal_expr) { - node.data.sort_properties = SortProperties::Ordered(options); - } - Ok(Transformed::yes(node)) -} - /// This function examines whether a referring expression directly refers to a /// given referred expression or if any of its children in the expression tree /// refer to the specified expression. diff --git a/datafusion/physical-expr/src/equivalence/properties/walk.rs b/datafusion/physical-expr/src/equivalence/properties/walk.rs new file mode 100644 index 0000000000000..a6cab2ae39dc6 --- /dev/null +++ b/datafusion/physical-expr/src/equivalence/properties/walk.rs @@ -0,0 +1,560 @@ +// 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 + +use std::sync::Arc; + +use super::EquivalenceProperties; +use crate::equivalence::ordering::is_partial_const; +use crate::expressions::{Column, Literal}; +use crate::{LexOrdering, PhysicalExpr, PhysicalSortExpr}; + +use arrow::compute::SortOptions; +use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; +use datafusion_common::{HashSet, Result}; +use datafusion_expr::interval_arithmetic::Interval; +use datafusion_expr::sort_properties::{ExprProperties, SortProperties}; +use datafusion_physical_expr_common::sort_expr::options_compatible; +use datafusion_physical_expr_common::utils::ExprPropertiesNode; + +use itertools::Itertools; + +/// The state of one walk along a sequence of sort keys, such as the keys of a +/// required ordering. +/// +/// A sort key only matters among the rows that tie on the keys before it. So +/// once a key is satisfied, the walk *pins* it: for the rest of the walk, the +/// key has one value, like a constant. An expression is *fixed* if it has one +/// value among those rows: it is a literal, a constant, pinned, a leaf that +/// reports itself as a singleton (such as an uncorrelated scalar subquery), or +/// not volatile and computed only from fixed expressions. The *head* of an +/// ordering is its first expression that is not fixed; an ordering whose +/// expressions are all fixed has no head. With orderings `[a, b]` and +/// `[a, c]`, the requirement `[a, c]` is satisfied, because pinning `a` makes +/// `c` the head of the second ordering. +/// +/// Pins live in the walk; the equivalence properties it reads do not change. +pub(super) struct SortKeyWalk<'a> { + props: &'a EquivalenceProperties, + /// Pinned expressions. Pinning a member of an equivalence class pins all + /// of its members. + pinned: HashSet>, + /// Orderings that hold among the rows that tie on the pinned expressions, + /// found by [`EquivalenceProperties::implied_ordering`] as the walk goes. + implied: Vec, + /// For each ordering (the normal orderings of `props`, then `implied`), + /// the index of its head. Empty until the first pin, so that a walk with + /// nothing pinned allocates nothing; until then every head is at index 0. + head_idxs: Vec, +} + +impl<'a> SortKeyWalk<'a> { + pub(super) fn new(props: &'a EquivalenceProperties) -> Self { + Self { + props, + pinned: HashSet::new(), + implied: vec![], + head_idxs: vec![], + } + } + + /// Returns how many leading keys (in normal form) the existing orderings + /// satisfy, pinning each satisfied key for the keys after it. A key + /// without sort options only needs to be ordered in some direction. + pub(super) fn sorted_prefix_len<'k>( + mut self, + keys: impl ExactSizeIterator, Option)>, + ) -> Result { + let schema = self.props.schema(); + let len = keys.len(); + for (idx, (expr, options)) in keys.enumerate() { + let satisfy = match self.expr_properties(Arc::clone(expr)).sort_properties { + SortProperties::Ordered(existing) => options.is_none_or(|options| { + let nullable = expr.nullable(schema).unwrap_or(true); + // `existing` comes from the first head that `expr` (or, + // for a computed `expr`, its arguments) matches; another + // head may be `expr` itself in the required direction. + options_compatible(&existing, &options, nullable) + || self.heads().any(|head| { + options_compatible(&head.options, &options, nullable) + && head.expr.eq(expr) + }) + }), + // Singleton expressions satisfy any ordering. + SortProperties::Singleton => true, + SortProperties::Unordered => false, + }; + if !satisfy { + return Ok(idx); + } + if idx + 1 < len { + self.pin_key(Arc::clone(expr))?; + } + } + Ok(len) + } + + /// Pins a satisfied sort key, and the sub-expressions it pins down: if an + /// expression is strictly order-preserving, equal outputs imply equal + /// values of its ordered children, so within a tie group of the key those + /// children have one value as well. For example, if data is sorted by + /// `[a, b]`, the requirement `[CAST(a AS BIGINT) ASC, b ASC]` is + /// satisfied: `a` has one value within each group of equal + /// `CAST(a AS BIGINT)` values, and hence `b` is sorted within each such + /// group. + fn pin_key(&mut self, key: Arc) -> Result<()> { + let mut stack = vec![key]; + while let Some(expr) = stack.pop() { + let properties = self.expr_properties(Arc::clone(&expr)); + if properties.strictly_order_preserving { + for child in expr.children() { + let child_properties = self.expr_properties(Arc::clone(child)); + if matches!( + child_properties.sort_properties, + SortProperties::Ordered(_) + ) { + stack.push(Arc::clone(child)); + } + } + } + self.pin(expr)?; + } + Ok(()) + } + + /// Pins `expr` and the rest of its equivalence class (but, unlike + /// [`Self::pin_key`], not its children), then moves the head of each + /// ordering past the expressions that are now fixed, and adds the + /// orderings that the heads imply, until no new ordering is found. + pub(super) fn pin(&mut self, expr: Arc) -> Result<()> { + match self.props.eq_group.get_equivalence_class(&expr) { + Some(class) => self.pinned.extend(class.iter().cloned()), + None => { + self.pinned.insert(expr); + } + } + let ordering_count = self.props.oeq_cache.normal_cls.len() + self.implied.len(); + self.head_idxs.resize(ordering_count, 0); + // Pins do not change below, so after the first round only the + // orderings found in the previous round can imply new ones. + let mut first = 0; + loop { + let count = self.head_idxs.len(); + let mut found = vec![]; + for idx in first..count { + let ordering = self.ordering(idx); + let mut head_idx = self.head_idxs[idx]; + while head_idx < ordering.len() && self.is_fixed(&ordering[head_idx].expr) + { + head_idx += 1; + } + if let Some(head) = ordering.get(head_idx) { + let rest = ordering[head_idx + 1..] + .iter() + .filter(|sort_expr| !self.is_fixed(&sort_expr.expr)); + found.extend(self.props.implied_ordering(head, rest)?); + } + self.head_idxs[idx] = head_idx; + } + for ordering in found.into_iter().filter_map(LexOrdering::new) { + let known = + (0..self.head_idxs.len()).any(|idx| *self.ordering(idx) == ordering); + if !known { + self.implied.push(ordering); + self.head_idxs.push(0); + } + } + // No new ordering: + if self.head_idxs.len() == count { + return Ok(()); + } + first = count; + } + } + + /// Returns the properties of `expr` among the rows that tie on the + /// pinned expressions. + pub(super) fn expr_properties(&self, expr: Arc) -> ExprProperties { + ExprPropertiesNode::new_unknown(expr) + .transform_up(|node| self.update_properties(node)) + .data() + .map(|node| node.data) + .unwrap_or_else(|_| ExprProperties::new_unknown()) + } + + /// Calculates the properties of a given [`ExprPropertiesNode`], whose + /// children are already done. + /// + /// - An intermediate node takes its order and range from its children, + /// through [`PhysicalExpr::get_properties`]. A `Literal` leaf is a + /// singleton with a point range. A `Column` leaf is unordered, with an + /// unbounded range of its data type. + /// - Then, whatever its children said: a node that is fixed or partially + /// constant is a singleton, and a node that heads an ordering is ordered + /// like the first ordering it heads. + fn update_properties( + &self, + mut node: ExprPropertiesNode, + ) -> Result> { + // First, try to gather the information from the children: + if !node.expr.children().is_empty() { + // We have an intermediate (non-leaf) node, account for its children: + let children_props = + node.children.iter().map(|c| c.data.clone()).collect_vec(); + // Properties that cannot be derived from the children, for example + // of a function whose input may be outside its domain, are + // unknown. The ordering of the node itself can still be known, so + // do not fail here, and keep the type of its range so that parents + // can derive from it. + node.data = node.expr.get_properties(&children_props).or_else(|_| { + Interval::make_unbounded(&node.expr.data_type(self.props.schema())?) + .map(|range| ExprProperties::new_unknown().with_range(range)) + })?; + } else if node.expr.is::() { + // We have a Literal, which is one of the two possible leaf node types: + node.data = node.expr.get_properties(&[])?; + } else if node.expr.is::() { + // We have a Column, which is the other possible leaf node type: + node.data.range = + Interval::make_unbounded(&node.expr.data_type(self.props.schema())?)?; + // A column is the identity mapping of itself, which is trivially + // strict: + node.data.strictly_order_preserving = true; + } + // Now, check what we know about orderings. A whole-partition window + // such as `SUM(x) OVER (PARTITION BY k)` adds `[k, w]` in all four + // directions of `w` (see `OrderingEquivalenceClass::is_expr_partial_const`). + // Once `k` is pinned, `w` heads orderings in opposite directions, so it + // has one value among the rows that tie on `k`: + let normal_expr = self.props.eq_group.normalize_expr(Arc::clone(&node.expr)); + if self.is_fixed(&normal_expr) || is_partial_const(self.heads(), &normal_expr) { + node.data.sort_properties = SortProperties::Singleton; + } else if let Some(head) = self.heads().find(|head| head.expr.eq(&normal_expr)) { + node.data.sort_properties = SortProperties::Ordered(head.options); + } + Ok(Transformed::yes(node)) + } + + /// Whether `expr` is fixed (see [`SortKeyWalk`]). Callers pass `expr` in + /// normal form. Pins are stored as given, so a pinned expression that is + /// not in normal form is found only through its equivalence class, whose + /// members are all pinned. + fn is_fixed(&self, expr: &Arc) -> bool { + if expr.is::() || self.pinned.contains(expr) { + return true; + } + let class = self.props.eq_group.get_equivalence_class(expr); + if class.is_some_and(|class| class.constant.is_some()) { + return true; + } + // A volatile expression can change from row to row, even when all of + // its inputs are fixed: + if expr.is_volatile_node() { + return false; + } + let children = expr.children(); + if children.is_empty() { + // A leaf that reads no column, such as an uncorrelated scalar + // subquery, is fixed if it reports itself as a singleton: + return expr + .get_properties(&[]) + .is_ok_and(|props| props.sort_properties == SortProperties::Singleton); + } + children.into_iter().all(|child| self.is_fixed(child)) + } + + /// The ordering at `idx`: the normal orderings of `props`, then + /// `implied`. + fn ordering(&self, idx: usize) -> &LexOrdering { + let normal = &self.props.oeq_cache.normal_cls; + normal + .get(idx) + .unwrap_or_else(|| &self.implied[idx - normal.len()]) + } + + /// The head of each ordering that has one. + fn heads(&self) -> impl Iterator + Clone { + let orderings = self.props.oeq_cache.normal_cls.iter().chain(&self.implied); + orderings.enumerate().filter_map(|(idx, ordering)| { + ordering.get(self.head_idxs.get(idx).copied().unwrap_or(0)) + }) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use crate::expressions::{BinaryExpr, CastExpr, col, lit}; + use crate::scalar_subquery::ScalarSubqueryExpr; + use crate::{AcrossPartitions, ConstExpr, EquivalenceProperties}; + use crate::{PhysicalExpr, PhysicalSortExpr, ScalarFunctionExpr}; + + use arrow::compute::SortOptions; + use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; + use datafusion_common::config::ConfigOptions; + use datafusion_common::{Result, ScalarValue}; + use datafusion_expr::physical_planning_context::{ + ScalarSubqueryResults, SubqueryIndex, + }; + use datafusion_expr::{ColumnarValue, Operator, Volatility, create_udf}; + + /// Pinning a key can move the head of an ordering to an expression that + /// is equal to a function of the expressions after it, which implies a + /// new ordering. With ordering `[k, a, c]` and `CAST(c AS DATE) = a`, + /// pinning `k` leaves `[a, c]`, which implies `[c]`, so `[k, c]` holds. + #[test] + fn ordering_found_after_pinning() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int32, true), + Field::new("a", DataType::Date32, true), + Field::new("c", DataType::Timestamp(TimeUnit::Nanosecond, None), true), + ])); + let [k, a, c] = ["k", "a", "c"].map(|name| col(name, &schema).unwrap()); + let cast_c = Arc::new(CastExpr::new_with_target_field( + Arc::clone(&c), + Arc::new(Field::new("c", DataType::Date32, true)), + None, + )) as _; + + let mut properties = EquivalenceProperties::new(Arc::clone(&schema)); + properties.add_ordering( + [&k, &a, &c].map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))), + ); + properties.add_equal_conditions(cast_c, Arc::clone(&a))?; + + let required = + [&k, &c].map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))); + assert!(properties.ordering_satisfy(required)?); + // Without `k` first, nothing is known about `c`: + assert!(!properties.ordering_satisfy([PhysicalSortExpr::new_default(c)])?); + Ok(()) + } + + /// An ordering found after a pin can imply another one. With ordering + /// `[k, x, y, z]`, `x = CAST(y AS BIGINT)` and `y = CAST(z AS INT)`, + /// pinning `k` leaves `[x, y, z]`, which implies `[y, z]`, which implies + /// `[z]`, so `[k, z]` holds. + #[test] + fn chained_orderings_found_after_pinning() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int32, true), + Field::new("x", DataType::Int64, true), + Field::new("y", DataType::Int32, true), + Field::new("z", DataType::Int16, true), + ])); + let [k, x, y, z] = ["k", "x", "y", "z"].map(|name| col(name, &schema).unwrap()); + let cast = |expr: &Arc, data_type| { + Arc::new(CastExpr::new_with_target_field( + Arc::clone(expr), + Arc::new(Field::new("cast", data_type, true)), + None, + )) as Arc + }; + + let mut properties = EquivalenceProperties::new(Arc::clone(&schema)); + properties.add_equal_conditions(Arc::clone(&y), cast(&z, DataType::Int32))?; + properties.add_equal_conditions(Arc::clone(&x), cast(&y, DataType::Int64))?; + properties.add_ordering( + [&k, &x, &y, &z].map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))), + ); + + let required = + [&k, &z].map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))); + assert!(properties.ordering_satisfy(required)?); + Ok(()) + } + + /// After pinning `k`, `a` heads three orderings, in two directions that + /// are not opposite. The requirement `[k, a ASC NULLS FIRST]` holds by the + /// last ordering, even though the first one heads `a` in another + /// direction. + #[test] + fn any_head_in_required_direction() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("k", DataType::Int32, true), + Field::new("c", DataType::Int32, true), + ])); + let [a, k, c] = ["a", "k", "c"].map(|name| col(name, &schema).unwrap()); + let sort = |expr: &Arc, descending| { + PhysicalSortExpr::new( + Arc::clone(expr), + SortOptions { + descending, + nulls_first: true, + }, + ) + }; + + let mut properties = EquivalenceProperties::new(Arc::clone(&schema)); + properties.add_orderings([ + vec![sort(&a, true), sort(&k, false), sort(&c, false)], + vec![sort(&k, false), sort(&a, true), sort(&c, false)], + vec![sort(&k, false), sort(&a, false)], + ]); + + assert!(properties.ordering_satisfy([sort(&k, false), sort(&a, false)])?); + Ok(()) + } + + /// A window over a whole partition, such as `SUM(..) OVER (PARTITION BY + /// k)`, adds `[k, x]` in all four directions of `x`. Pinning `k` leaves `x` + /// heading orderings in opposite directions, so `x` has one value, and so + /// does `CAST(x AS BIGINT)`: `[k, CAST(x AS BIGINT)]` holds in every + /// direction. Without the partial constant check, `x` would only be + /// ordered in the direction of the first ordering it heads. + #[test] + fn partial_constant_after_pinning() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int32, true), + Field::new("x", DataType::Int32, true), + ])); + let [k, x] = ["k", "x"].map(|name| col(name, &schema).unwrap()); + let cast_x = Arc::new(CastExpr::new_with_target_field( + Arc::clone(&x), + Arc::new(Field::new("x", DataType::Int64, true)), + None, + )) as Arc; + let directions = [(false, false), (false, true), (true, false), (true, true)] + .map(|(descending, nulls_first)| SortOptions { + descending, + nulls_first, + }); + + let mut properties = EquivalenceProperties::new(Arc::clone(&schema)); + properties.add_orderings(directions.map(|options| { + vec![ + PhysicalSortExpr::new_default(Arc::clone(&k)), + PhysicalSortExpr::new(Arc::clone(&x), options), + ] + })); + + for options in directions { + let required = [ + PhysicalSortExpr::new_default(Arc::clone(&k)), + PhysicalSortExpr::new(Arc::clone(&cast_x), options), + ]; + assert!(properties.ordering_satisfy(required)?, "{options:?}"); + } + Ok(()) + } + + /// `find_longest_permutation` pins what it finds in normal form. With + /// `b = c` and ordering `[a + b, d]`, finding either `a + b` or `a + c` + /// pins the head of the ordering, whichever of `b` and `c` the class + /// normalizes to, so `d` is found next. + #[test] + fn find_longest_permutation_pins_normal_form() -> Result<()> { + let schema = Arc::new(Schema::new( + ["a", "b", "c", "d"] + .map(|name| Field::new(name, DataType::Int32, true)) + .to_vec(), + )); + let [a, b, c, d] = ["a", "b", "c", "d"].map(|name| col(name, &schema).unwrap()); + let plus = |right: &Arc| { + Arc::new(BinaryExpr::new( + Arc::clone(&a), + Operator::Plus, + Arc::clone(right), + )) as Arc + }; + + let mut properties = EquivalenceProperties::new(Arc::clone(&schema)); + properties.add_equal_conditions(Arc::clone(&b), Arc::clone(&c))?; + properties.add_ordering([ + PhysicalSortExpr::new_default(plus(&b)), + PhysicalSortExpr::new_default(Arc::clone(&d)), + ]); + + for right in [&b, &c] { + let (_, indices) = + properties.find_longest_permutation(&[plus(right), Arc::clone(&d)])?; + assert_eq!(indices, vec![0, 1], "a + {right}"); + } + Ok(()) + } + + /// A literal key is found as a singleton and does not change the + /// properties. With `x` the constant `5` and ordering `[k, b % 5]`, + /// `[5, k, b % 5]` is found in full. (Adding `5` as a constant, as the + /// properties were changed before, put it in the class of `x`, so `b % 5` + /// no longer matched the ordering.) + #[test] + fn find_longest_permutation_with_literal() -> Result<()> { + let schema = Arc::new(Schema::new( + ["x", "k", "b"] + .map(|name| Field::new(name, DataType::Int32, true)) + .to_vec(), + )); + let [x, k, b] = ["x", "k", "b"].map(|name| col(name, &schema).unwrap()); + let b_mod_5 = Arc::new(BinaryExpr::new(b, Operator::Modulo, lit(5i32))) as _; + + let mut properties = EquivalenceProperties::new(Arc::clone(&schema)); + properties.add_constants([ConstExpr::new( + x, + AcrossPartitions::Uniform(Some(ScalarValue::Int32(Some(5)))), + )])?; + properties.add_ordering([ + PhysicalSortExpr::new_default(Arc::clone(&k)), + PhysicalSortExpr::new_default(Arc::clone(&b_mod_5)), + ]); + + let (_, indices) = + properties.find_longest_permutation(&[lit(5i32), k, b_mod_5])?; + assert_eq!(indices, vec![0, 1, 2]); + Ok(()) + } + + /// A volatile function can return a new value for every row, so it is not + /// fixed even when its argument is pinned. An uncorrelated scalar subquery + /// has one value everywhere, so it is fixed, and so is `a + subquery` once + /// `a` is pinned. (A requirement drops a key that is just the subquery + /// before the walk sees it.) + #[test] + fn volatile_function_and_scalar_subquery_keys() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let a = col("a", &schema)?; + let subquery = Arc::new(ScalarSubqueryExpr::new( + DataType::Int32, + true, + SubqueryIndex::new(0), + ScalarSubqueryResults::new(1), + )) as Arc; + let call = |volatility| -> Result> { + let udf = Arc::new(create_udf( + "f", + vec![DataType::Int32], + DataType::Int32, + volatility, + Arc::new(|args: &[ColumnarValue]| Ok(args[0].clone())), + )); + Ok(Arc::new(ScalarFunctionExpr::try_new( + udf, + vec![Arc::clone(&a)], + &schema, + Arc::new(ConfigOptions::new()), + )?)) + }; + + let mut properties = EquivalenceProperties::new(Arc::clone(&schema)); + properties.add_ordering([PhysicalSortExpr::new_default(Arc::clone(&a))]); + let satisfied = |exprs: [&Arc; 2]| { + properties.ordering_satisfy( + exprs.map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))), + ) + }; + + assert!(!satisfied([&a, &call(Volatility::Volatile)?])?); + assert!(satisfied([&a, &call(Volatility::Stable)?])?); + let a_plus_subquery = Arc::new(BinaryExpr::new( + Arc::clone(&a), + Operator::Plus, + Arc::clone(&subquery), + )) as Arc; + assert!(satisfied([&a, &a_plus_subquery])?); + let (_, found) = properties.find_longest_permutation(&[subquery])?; + assert_eq!(found, [0]); + Ok(()) + } +}