use std::fmt::Display;
use std::hash::Hash;
use std::sync::Arc;
use std::vec::IntoIter;
use crate::equivalence::add_offset_to_expr;
use crate::{LexOrdering, PhysicalExpr};
use arrow_schema::SortOptions;
#[derive(Debug, Clone, Eq, PartialEq, Hash, Default)]
pub struct OrderingEquivalenceClass {
pub orderings: Vec<LexOrdering>,
}
impl OrderingEquivalenceClass {
pub fn empty() -> Self {
Default::default()
}
pub fn clear(&mut self) {
self.orderings.clear();
}
pub fn new(orderings: Vec<LexOrdering>) -> Self {
let mut result = Self { orderings };
result.remove_redundant_entries();
result
}
pub fn contains(&self, ordering: &LexOrdering) -> bool {
self.orderings.contains(ordering)
}
#[allow(dead_code)]
fn push(&mut self, ordering: LexOrdering) {
self.orderings.push(ordering);
self.remove_redundant_entries();
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn iter(&self) -> impl Iterator<Item = &LexOrdering> {
self.orderings.iter()
}
pub fn len(&self) -> usize {
self.orderings.len()
}
pub fn extend(&mut self, other: Self) {
self.orderings.extend(other.orderings);
self.remove_redundant_entries();
}
pub fn add_new_orderings(
&mut self,
orderings: impl IntoIterator<Item = LexOrdering>,
) {
self.orderings.extend(orderings);
self.remove_redundant_entries();
}
pub fn add_new_ordering(&mut self, ordering: LexOrdering) {
self.add_new_orderings([ordering]);
}
fn remove_redundant_entries(&mut self) {
let mut work = true;
while work {
work = false;
let mut idx = 0;
while idx < self.orderings.len() {
let mut ordering_idx = idx + 1;
let mut removal = self.orderings[idx].is_empty();
while ordering_idx < self.orderings.len() {
work |= resolve_overlap(&mut self.orderings, idx, ordering_idx);
if self.orderings[idx].is_empty() {
removal = true;
break;
}
work |= resolve_overlap(&mut self.orderings, ordering_idx, idx);
if self.orderings[ordering_idx].is_empty() {
self.orderings.swap_remove(ordering_idx);
} else {
ordering_idx += 1;
}
}
if removal {
self.orderings.swap_remove(idx);
} else {
idx += 1;
}
}
}
}
pub fn output_ordering(&self) -> Option<LexOrdering> {
let output_ordering = self
.orderings
.iter()
.flat_map(|ordering| ordering.as_ref())
.cloned()
.collect();
let output_ordering = collapse_lex_ordering(output_ordering);
(!output_ordering.is_empty()).then_some(output_ordering)
}
pub fn join_suffix(mut self, other: &Self) -> Self {
let n_ordering = self.orderings.len();
let n_cross = std::cmp::max(n_ordering, other.len() * n_ordering);
self.orderings = self
.orderings
.iter()
.cloned()
.cycle()
.take(n_cross)
.collect();
for (outer_idx, ordering) in other.iter().enumerate() {
for idx in 0..n_ordering {
let idx = outer_idx * n_ordering + idx;
self.orderings[idx].inner.extend(ordering.iter().cloned());
}
}
self
}
pub fn add_offset(&mut self, offset: usize) {
for ordering in self.orderings.iter_mut() {
for sort_expr in ordering.inner.iter_mut() {
sort_expr.expr = add_offset_to_expr(Arc::clone(&sort_expr.expr), offset);
}
}
}
pub fn get_options(&self, expr: &Arc<dyn PhysicalExpr>) -> Option<SortOptions> {
for ordering in self.iter() {
let leading_ordering = &ordering[0];
if leading_ordering.expr.eq(expr) {
return Some(leading_ordering.options);
}
}
None
}
}
impl IntoIterator for OrderingEquivalenceClass {
type Item = LexOrdering;
type IntoIter = IntoIter<LexOrdering>;
fn into_iter(self) -> Self::IntoIter {
self.orderings.into_iter()
}
}
pub fn collapse_lex_ordering(input: LexOrdering) -> LexOrdering {
let mut output = LexOrdering::default();
for item in input.iter() {
if !output.iter().any(|req| req.expr.eq(&item.expr)) {
output.push(item.clone());
}
}
output
}
fn resolve_overlap(orderings: &mut [LexOrdering], idx: usize, pre_idx: usize) -> bool {
let length = orderings[idx].len();
let other_length = orderings[pre_idx].len();
for overlap in 1..=length.min(other_length) {
if orderings[idx][length - overlap..] == orderings[pre_idx][..overlap] {
orderings[idx].truncate(length - overlap);
return true;
}
}
false
}
impl Display for OrderingEquivalenceClass {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "[")?;
let mut iter = self.orderings.iter();
if let Some(ordering) = iter.next() {
write!(f, "[{}]", ordering)?;
}
for ordering in iter {
write!(f, ", [{}]", ordering)?;
}
write!(f, "]")?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use crate::equivalence::tests::{
convert_to_orderings, convert_to_sort_exprs, create_test_schema,
};
use crate::equivalence::{
EquivalenceClass, EquivalenceGroup, EquivalenceProperties,
OrderingEquivalenceClass,
};
use crate::expressions::{col, BinaryExpr, Column};
use crate::utils::tests::TestScalarUDF;
use crate::{ConstExpr, PhysicalExpr, PhysicalSortExpr};
use arrow::datatypes::{DataType, Field, Schema};
use arrow_schema::SortOptions;
use datafusion_common::{DFSchema, Result};
use datafusion_expr::{Operator, ScalarUDF};
use datafusion_physical_expr_common::sort_expr::LexOrdering;
#[test]
fn test_ordering_satisfy() -> Result<()> {
let input_schema = Arc::new(Schema::new(vec![
Field::new("a", DataType::Int64, true),
Field::new("b", DataType::Int64, true),
]));
let crude = LexOrdering::new(vec![PhysicalSortExpr {
expr: Arc::new(Column::new("a", 0)),
options: SortOptions::default(),
}]);
let finer = LexOrdering::new(vec![
PhysicalSortExpr {
expr: Arc::new(Column::new("a", 0)),
options: SortOptions::default(),
},
PhysicalSortExpr {
expr: Arc::new(Column::new("b", 1)),
options: SortOptions::default(),
},
]);
let mut eq_properties_finer =
EquivalenceProperties::new(Arc::clone(&input_schema));
eq_properties_finer.oeq_class.push(finer.clone());
assert!(eq_properties_finer.ordering_satisfy(crude.as_ref()));
let mut eq_properties_crude =
EquivalenceProperties::new(Arc::clone(&input_schema));
eq_properties_crude.oeq_class.push(crude);
assert!(!eq_properties_crude.ordering_satisfy(finer.as_ref()));
Ok(())
}
#[test]
fn test_ordering_satisfy_with_equivalence2() -> Result<()> {
let test_schema = create_test_schema()?;
let col_a = &col("a", &test_schema)?;
let col_b = &col("b", &test_schema)?;
let col_c = &col("c", &test_schema)?;
let col_d = &col("d", &test_schema)?;
let col_e = &col("e", &test_schema)?;
let col_f = &col("f", &test_schema)?;
let test_fun = ScalarUDF::new_from_impl(TestScalarUDF::new());
let floor_a = &crate::udf::create_physical_expr(
&test_fun,
&[col("a", &test_schema)?],
&test_schema,
&[],
&DFSchema::empty(),
)?;
let floor_f = &crate::udf::create_physical_expr(
&test_fun,
&[col("f", &test_schema)?],
&test_schema,
&[],
&DFSchema::empty(),
)?;
let exp_a = &crate::udf::create_physical_expr(
&test_fun,
&[col("a", &test_schema)?],
&test_schema,
&[],
&DFSchema::empty(),
)?;
let a_plus_b = Arc::new(BinaryExpr::new(
Arc::clone(col_a),
Operator::Plus,
Arc::clone(col_b),
)) as Arc<dyn PhysicalExpr>;
let options = SortOptions {
descending: false,
nulls_first: false,
};
let test_cases = vec![
(
vec![
vec![(col_a, options), (col_d, options), (col_b, options)],
vec![(col_c, options)],
],
vec![vec![col_a, col_f]],
vec![col_e],
vec![(col_a, options), (col_b, options)],
false,
),
(
vec![
vec![(col_a, options), (col_c, options), (col_b, options)],
vec![(col_d, options)],
],
vec![vec![col_a, col_f]],
vec![col_e],
vec![(floor_a, options)],
true,
),
(
vec![
vec![(col_a, options), (col_c, options), (col_b, options)],
vec![(col_d, options)],
],
vec![vec![col_a, col_f]],
vec![col_e],
vec![(floor_f, options)],
true,
),
(
vec![
vec![(col_a, options), (col_c, options), (col_b, options)],
vec![(col_d, options)],
],
vec![vec![col_a, col_f]],
vec![col_e],
vec![(col_a, options), (col_c, options), (&a_plus_b, options)],
true,
),
(
vec![
vec![
(col_a, options),
(col_b, options),
(col_c, options),
(col_d, options),
],
],
vec![vec![col_a, col_f]],
vec![col_e],
vec![(floor_a, options), (&a_plus_b, options)],
false,
),
(
vec![
vec![
(col_a, options),
(col_b, options),
(col_c, options),
(col_d, options),
],
],
vec![vec![col_a, col_f]],
vec![col_e],
vec![(exp_a, options), (&a_plus_b, options)],
false,
),
(
vec![
vec![(col_a, options), (col_d, options), (col_b, options)],
vec![(col_c, options)],
],
vec![vec![col_a, col_f]],
vec![col_e],
vec![(col_a, options), (col_d, options), (floor_a, options)],
true,
),
(
vec![
vec![(col_a, options), (col_c, options), (col_b, options)],
vec![(col_d, options)],
],
vec![vec![col_a, col_f]],
vec![col_e],
vec![(col_a, options), (floor_a, options), (&a_plus_b, options)],
false,
),
(
vec![
vec![(col_a, options), (col_b, options), (col_c, options)],
vec![(col_d, options)],
],
vec![vec![col_a, col_f]],
vec![col_e],
vec![
(col_a, options),
(col_c, options),
(floor_a, options),
(&a_plus_b, options),
],
false,
),
(
vec![
vec![
(col_a, options),
(col_b, options),
(col_c, options),
(col_d, options),
],
],
vec![vec![col_a, col_f]],
vec![col_e],
vec![
(col_a, options),
(col_b, options),
(col_c, options),
(floor_a, options),
],
true,
),
(
vec![
vec![(col_d, options), (col_b, options)],
vec![(col_c, options), (col_a, options)],
],
vec![vec![col_a, col_f]],
vec![col_e],
vec![(col_c, options), (col_d, options), (&a_plus_b, options)],
true,
),
];
for (orderings, eq_group, constants, reqs, expected) in test_cases {
let err_msg =
format!("error in test orderings: {orderings:?}, eq_group: {eq_group:?}, constants: {constants:?}, reqs: {reqs:?}, expected: {expected:?}");
let mut eq_properties = EquivalenceProperties::new(Arc::clone(&test_schema));
let orderings = convert_to_orderings(&orderings);
eq_properties.add_new_orderings(orderings);
let eq_group = eq_group
.into_iter()
.map(|eq_class| {
let eq_classes = eq_class.into_iter().cloned().collect::<Vec<_>>();
EquivalenceClass::new(eq_classes)
})
.collect::<Vec<_>>();
let eq_group = EquivalenceGroup::new(eq_group);
eq_properties.add_equivalence_group(eq_group);
let constants = constants
.into_iter()
.map(|expr| ConstExpr::from(expr).with_across_partitions(true));
eq_properties = eq_properties.with_constants(constants);
let reqs = convert_to_sort_exprs(&reqs);
assert_eq!(
eq_properties.ordering_satisfy(reqs.as_ref()),
expected,
"{}",
err_msg
);
}
Ok(())
}
#[test]
fn test_ordering_satisfy_different_lengths() -> Result<()> {
let test_schema = create_test_schema()?;
let col_a = &col("a", &test_schema)?;
let col_b = &col("b", &test_schema)?;
let col_c = &col("c", &test_schema)?;
let col_d = &col("d", &test_schema)?;
let col_e = &col("e", &test_schema)?;
let col_f = &col("f", &test_schema)?;
let options = SortOptions {
descending: false,
nulls_first: false,
};
let mut eq_properties = EquivalenceProperties::new(test_schema);
eq_properties.add_equal_conditions(col_a, col_c)?;
let orderings = vec![
vec![(col_a, options)],
vec![(col_e, options)],
vec![(col_d, options), (col_f, options)],
];
let orderings = convert_to_orderings(&orderings);
eq_properties.add_new_orderings(orderings);
let test_cases = vec![
(
vec![(col_c, options), (col_a, options), (col_e, options)],
true,
),
(vec![(col_c, options), (col_b, options)], false),
(vec![(col_c, options), (col_d, options)], true),
(
vec![(col_d, options), (col_f, options), (col_b, options)],
false,
),
(vec![(col_d, options), (col_f, options)], true),
];
for (reqs, expected) in test_cases {
let err_msg =
format!("error in test reqs: {:?}, expected: {:?}", reqs, expected,);
let reqs = convert_to_sort_exprs(&reqs);
assert_eq!(
eq_properties.ordering_satisfy(reqs.as_ref()),
expected,
"{}",
err_msg
);
}
Ok(())
}
#[test]
fn test_remove_redundant_entries_oeq_class() -> Result<()> {
let schema = create_test_schema()?;
let col_a = &col("a", &schema)?;
let col_b = &col("b", &schema)?;
let col_c = &col("c", &schema)?;
let col_d = &col("d", &schema)?;
let col_e = &col("e", &schema)?;
let option_asc = SortOptions {
descending: false,
nulls_first: false,
};
let option_desc = SortOptions {
descending: true,
nulls_first: true,
};
let test_cases = vec![
(
vec![
vec![(col_a, option_asc), (col_b, option_asc)],
],
vec![
vec![(col_a, option_asc), (col_b, option_asc)],
],
),
(
vec![
vec![(col_a, option_asc), (col_b, option_asc)],
vec![
(col_a, option_asc),
(col_b, option_asc),
(col_c, option_asc),
],
],
vec![
vec![
(col_a, option_asc),
(col_b, option_asc),
(col_c, option_asc),
],
],
),
(
vec![
vec![(col_a, option_asc), (col_b, option_desc)],
vec![(col_a, option_asc)],
vec![(col_a, option_asc), (col_c, option_asc)],
],
vec![
vec![(col_a, option_asc), (col_b, option_desc)],
vec![(col_a, option_asc), (col_c, option_asc)],
],
),
(
vec![
vec![(col_a, option_asc), (col_b, option_asc)],
vec![
(col_a, option_asc),
(col_b, option_asc),
(col_c, option_asc),
],
vec![(col_a, option_asc)],
],
vec![
vec![
(col_a, option_asc),
(col_b, option_asc),
(col_c, option_asc),
],
],
),
(
vec![vec![]],
vec![],
),
(
vec![
vec![(col_a, option_asc), (col_b, option_asc)],
vec![(col_b, option_asc)],
],
vec![
vec![(col_a, option_asc)],
vec![(col_b, option_asc)],
],
),
(
vec![
vec![(col_b, option_asc), (col_a, option_asc)],
vec![(col_c, option_asc), (col_a, option_asc)],
vec![
(col_d, option_asc),
(col_b, option_asc),
(col_c, option_asc),
],
],
vec![
vec![(col_b, option_asc), (col_a, option_asc)],
vec![(col_c, option_asc), (col_a, option_asc)],
vec![(col_d, option_asc)],
],
),
(
vec![
vec![(col_b, option_asc), (col_e, option_asc)],
vec![(col_c, option_asc), (col_a, option_asc)],
vec![
(col_d, option_asc),
(col_b, option_asc),
(col_e, option_asc),
(col_c, option_asc),
(col_a, option_asc),
],
],
vec![
vec![(col_b, option_asc), (col_e, option_asc)],
vec![(col_c, option_asc), (col_a, option_asc)],
vec![(col_d, option_asc)],
],
),
(
vec![
vec![(col_b, option_asc)],
vec![
(col_a, option_asc),
(col_b, option_asc),
(col_c, option_asc),
],
vec![
(col_d, option_asc),
(col_a, option_asc),
(col_b, option_asc),
],
],
vec![
vec![(col_b, option_asc)],
vec![
(col_a, option_asc),
(col_b, option_asc),
(col_c, option_asc),
],
vec![(col_d, option_asc)],
],
),
];
for (orderings, expected) in test_cases {
let orderings = convert_to_orderings(&orderings);
let expected = convert_to_orderings(&expected);
let actual = OrderingEquivalenceClass::new(orderings.clone());
let actual = actual.orderings;
let err_msg = format!(
"orderings: {:?}, expected: {:?}, actual :{:?}",
orderings, expected, actual
);
assert_eq!(actual.len(), expected.len(), "{}", err_msg);
for elem in actual {
assert!(expected.contains(&elem), "{}", err_msg);
}
}
Ok(())
}
}