in datafusion/physical-expr/src/equivalence/ordering.rs [403:693]
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 = Arc::new(ScalarUDF::new_from_impl(TestScalarUDF::new()));
let floor_a = Arc::new(ScalarFunctionExpr::try_new(
Arc::clone(&test_fun),
vec![Arc::clone(col_a)],
&test_schema,
)?) as PhysicalExprRef;
let floor_f = Arc::new(ScalarFunctionExpr::try_new(
Arc::clone(&test_fun),
vec![Arc::clone(col_f)],
&test_schema,
)?) as PhysicalExprRef;
let exp_a = Arc::new(ScalarFunctionExpr::try_new(
Arc::clone(&test_fun),
vec![Arc::clone(col_a)],
&test_schema,
)?) as PhysicalExprRef;
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![
// ------------ TEST CASE 1 ------------
(
// orderings
vec![
// [a ASC, d ASC, b ASC]
vec![(col_a, options), (col_d, options), (col_b, options)],
// [c ASC]
vec![(col_c, options)],
],
// equivalence classes
vec![vec![col_a, col_f]],
// constants
vec![col_e],
// requirement [a ASC, b ASC], requirement is not satisfied.
vec![(col_a, options), (col_b, options)],
// expected: requirement is not satisfied.
false,
),
// ------------ TEST CASE 2 ------------
(
// orderings
vec![
// [a ASC, c ASC, b ASC]
vec![(col_a, options), (col_c, options), (col_b, options)],
// [d ASC]
vec![(col_d, options)],
],
// equivalence classes
vec![vec![col_a, col_f]],
// constants
vec![col_e],
// requirement [floor(a) ASC],
vec![(&floor_a, options)],
// expected: requirement is satisfied.
true,
),
// ------------ TEST CASE 2.1 ------------
(
// orderings
vec![
// [a ASC, c ASC, b ASC]
vec![(col_a, options), (col_c, options), (col_b, options)],
// [d ASC]
vec![(col_d, options)],
],
// equivalence classes
vec![vec![col_a, col_f]],
// constants
vec![col_e],
// requirement [floor(f) ASC], (Please note that a=f)
vec![(&floor_f, options)],
// expected: requirement is satisfied.
true,
),
// ------------ TEST CASE 3 ------------
(
// orderings
vec![
// [a ASC, c ASC, b ASC]
vec![(col_a, options), (col_c, options), (col_b, options)],
// [d ASC]
vec![(col_d, options)],
],
// equivalence classes
vec![vec![col_a, col_f]],
// constants
vec![col_e],
// requirement [a ASC, c ASC, a+b ASC],
vec![(col_a, options), (col_c, options), (&a_plus_b, options)],
// expected: requirement is satisfied.
true,
),
// ------------ TEST CASE 4 ------------
(
// orderings
vec![
// [a ASC, b ASC, c ASC, d ASC]
vec![
(col_a, options),
(col_b, options),
(col_c, options),
(col_d, options),
],
],
// equivalence classes
vec![vec![col_a, col_f]],
// constants
vec![col_e],
// requirement [floor(a) ASC, a+b ASC],
vec![(&floor_a, options), (&a_plus_b, options)],
// expected: requirement is satisfied.
false,
),
// ------------ TEST CASE 5 ------------
(
// orderings
vec![
// [a ASC, b ASC, c ASC, d ASC]
vec![
(col_a, options),
(col_b, options),
(col_c, options),
(col_d, options),
],
],
// equivalence classes
vec![vec![col_a, col_f]],
// constants
vec![col_e],
// requirement [exp(a) ASC, a+b ASC],
vec![(&exp_a, options), (&a_plus_b, options)],
// expected: requirement is not satisfied.
// TODO: If we know that exp function is 1-to-1 function.
// we could have deduced that above requirement is satisfied.
false,
),
// ------------ TEST CASE 6 ------------
(
// orderings
vec![
// [a ASC, d ASC, b ASC]
vec![(col_a, options), (col_d, options), (col_b, options)],
// [c ASC]
vec![(col_c, options)],
],
// equivalence classes
vec![vec![col_a, col_f]],
// constants
vec![col_e],
// requirement [a ASC, d ASC, floor(a) ASC],
vec![(col_a, options), (col_d, options), (&floor_a, options)],
// expected: requirement is satisfied.
true,
),
// ------------ TEST CASE 7 ------------
(
// orderings
vec![
// [a ASC, c ASC, b ASC]
vec![(col_a, options), (col_c, options), (col_b, options)],
// [d ASC]
vec![(col_d, options)],
],
// equivalence classes
vec![vec![col_a, col_f]],
// constants
vec![col_e],
// requirement [a ASC, floor(a) ASC, a + b ASC],
vec![(col_a, options), (&floor_a, options), (&a_plus_b, options)],
// expected: requirement is not satisfied.
false,
),
// ------------ TEST CASE 8 ------------
(
// orderings
vec![
// [a ASC, b ASC, c ASC]
vec![(col_a, options), (col_b, options), (col_c, options)],
// [d ASC]
vec![(col_d, options)],
],
// equivalence classes
vec![vec![col_a, col_f]],
// constants
vec![col_e],
// requirement [a ASC, c ASC, floor(a) ASC, a + b ASC],
vec![
(col_a, options),
(col_c, options),
(&floor_a, options),
(&a_plus_b, options),
],
// expected: requirement is not satisfied.
false,
),
// ------------ TEST CASE 9 ------------
(
// orderings
vec![
// [a ASC, b ASC, c ASC, d ASC]
vec![
(col_a, options),
(col_b, options),
(col_c, options),
(col_d, options),
],
],
// equivalence classes
vec![vec![col_a, col_f]],
// constants
vec![col_e],
// requirement [a ASC, b ASC, c ASC, floor(a) ASC],
vec![
(col_a, options),
(col_b, options),
(col_c, options),
(&floor_a, options),
],
// expected: requirement is satisfied.
true,
),
// ------------ TEST CASE 10 ------------
(
// orderings
vec![
// [d ASC, b ASC]
vec![(col_d, options), (col_b, options)],
// [c ASC, a ASC]
vec![(col_c, options), (col_a, options)],
],
// equivalence classes
vec![vec![col_a, col_f]],
// constants
vec![col_e],
// requirement [c ASC, d ASC, a + b ASC],
vec![(col_c, options), (col_d, options), (&a_plus_b, options)],
// expected: requirement is satisfied.
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(AcrossPartitions::Uniform(None))
});
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(())
}