use std::fmt;
use std::fmt::Display;
use std::hash::{Hash, Hasher};
use std::iter::Peekable;
use std::slice::Iter;
use std::sync::Arc;
use super::ordering::collapse_lex_ordering;
use crate::equivalence::class::const_exprs_contains;
use crate::equivalence::{
collapse_lex_req, EquivalenceClass, EquivalenceGroup, OrderingEquivalenceClass,
ProjectionMapping,
};
use crate::expressions::{with_new_schema, CastExpr, Column, Literal};
use crate::{
physical_exprs_contains, ConstExpr, LexOrdering, LexOrderingRef, LexRequirement,
LexRequirementRef, PhysicalExpr, PhysicalExprRef, PhysicalSortExpr,
PhysicalSortRequirement,
};
use arrow_schema::{SchemaRef, SortOptions};
use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
use datafusion_common::{internal_err, plan_err, JoinSide, JoinType, Result};
use datafusion_expr::interval_arithmetic::Interval;
use datafusion_expr::sort_properties::{ExprProperties, SortProperties};
use datafusion_physical_expr_common::utils::ExprPropertiesNode;
use indexmap::{IndexMap, IndexSet};
use itertools::Itertools;
#[derive(Debug, Clone)]
pub struct EquivalenceProperties {
pub eq_group: EquivalenceGroup,
pub oeq_class: OrderingEquivalenceClass,
pub constants: Vec<ConstExpr>,
schema: SchemaRef,
}
impl EquivalenceProperties {
pub fn new(schema: SchemaRef) -> Self {
Self {
eq_group: EquivalenceGroup::empty(),
oeq_class: OrderingEquivalenceClass::empty(),
constants: vec![],
schema,
}
}
pub fn new_with_orderings(schema: SchemaRef, orderings: &[LexOrdering]) -> Self {
Self {
eq_group: EquivalenceGroup::empty(),
oeq_class: OrderingEquivalenceClass::new(orderings.to_vec()),
constants: vec![],
schema,
}
}
pub fn schema(&self) -> &SchemaRef {
&self.schema
}
pub fn oeq_class(&self) -> &OrderingEquivalenceClass {
&self.oeq_class
}
pub fn eq_group(&self) -> &EquivalenceGroup {
&self.eq_group
}
pub fn constants(&self) -> &[ConstExpr] {
&self.constants
}
pub fn output_ordering(&self) -> Option<LexOrdering> {
let constants = self.constants();
let mut output_ordering = self.oeq_class().output_ordering().unwrap_or_default();
output_ordering
.inner
.retain(|sort_expr| !const_exprs_contains(constants, &sort_expr.expr));
(!output_ordering.is_empty()).then_some(output_ordering)
}
pub fn normalized_oeq_class(&self) -> OrderingEquivalenceClass {
OrderingEquivalenceClass::new(
self.oeq_class
.iter()
.map(|ordering| self.normalize_sort_exprs(ordering.as_ref()))
.collect(),
)
}
pub fn extend(mut self, other: Self) -> Self {
self.eq_group.extend(other.eq_group);
self.oeq_class.extend(other.oeq_class);
self.with_constants(other.constants)
}
pub fn clear_orderings(&mut self) {
self.oeq_class.clear();
}
pub fn clear_per_partition_constants(&mut self) {
self.constants.retain(|item| item.across_partitions());
}
pub fn add_ordering_equivalence_class(&mut self, other: OrderingEquivalenceClass) {
self.oeq_class.extend(other);
}
pub fn add_new_orderings(
&mut self,
orderings: impl IntoIterator<Item = LexOrdering>,
) {
self.oeq_class.add_new_orderings(orderings);
}
pub fn add_new_ordering(&mut self, ordering: LexOrdering) {
self.add_new_orderings([ordering]);
}
pub fn add_equivalence_group(&mut self, other_eq_group: EquivalenceGroup) {
self.eq_group.extend(other_eq_group);
}
pub fn add_equal_conditions(
&mut self,
left: &Arc<dyn PhysicalExpr>,
right: &Arc<dyn PhysicalExpr>,
) -> Result<()> {
if self.is_expr_constant(left) {
if !const_exprs_contains(&self.constants, right) {
self.constants
.push(ConstExpr::from(right).with_across_partitions(true));
}
} else if self.is_expr_constant(right) {
if !const_exprs_contains(&self.constants, left) {
self.constants
.push(ConstExpr::from(left).with_across_partitions(true));
}
}
self.eq_group.add_equal_conditions(left, right);
self.discover_new_orderings(left)?;
Ok(())
}
#[deprecated(since = "43.0.0", note = "Use [`with_constants`] instead")]
pub fn add_constants(self, constants: impl IntoIterator<Item = ConstExpr>) -> Self {
self.with_constants(constants)
}
pub fn remove_constant(mut self, c: &ConstExpr) -> Self {
self.constants.retain(|existing| existing != c);
self
}
pub fn with_constants(
mut self,
constants: impl IntoIterator<Item = ConstExpr>,
) -> Self {
let (const_exprs, across_partition_flags): (
Vec<Arc<dyn PhysicalExpr>>,
Vec<bool>,
) = constants
.into_iter()
.map(|const_expr| {
let across_partitions = const_expr.across_partitions();
let expr = const_expr.owned_expr();
(expr, across_partitions)
})
.unzip();
for (expr, across_partitions) in self
.eq_group
.normalize_exprs(const_exprs)
.into_iter()
.zip(across_partition_flags)
{
if !const_exprs_contains(&self.constants, &expr) {
let const_expr =
ConstExpr::from(expr).with_across_partitions(across_partitions);
self.constants.push(const_expr);
}
}
for ordering in self.normalized_oeq_class().iter() {
if let Err(e) = self.discover_new_orderings(&ordering[0].expr) {
log::debug!("error discovering new orderings: {e}");
}
}
self
}
fn discover_new_orderings(&mut self, expr: &Arc<dyn PhysicalExpr>) -> Result<()> {
let normalized_expr = self.eq_group().normalize_expr(Arc::clone(expr));
let eq_class = self
.eq_group
.classes
.iter()
.find_map(|class| {
class
.contains(&normalized_expr)
.then(|| class.clone().into_vec())
})
.unwrap_or_else(|| vec![Arc::clone(&normalized_expr)]);
let mut new_orderings: Vec<LexOrdering> = vec![];
for (ordering, next_expr) in self
.normalized_oeq_class()
.iter()
.filter(|ordering| ordering[0].expr.eq(&normalized_expr))
.filter_map(|ordering| Some(ordering).zip(ordering.inner.get(1)))
{
let leading_ordering = ordering[0].options;
for equivalent_expr in &eq_class {
let children = equivalent_expr.children();
if children.len() == 1
&& children[0].eq(&next_expr.expr)
&& SortProperties::Ordered(leading_ordering)
== equivalent_expr
.get_properties(&[ExprProperties {
sort_properties: SortProperties::Ordered(
leading_ordering,
),
range: Interval::make_unbounded(
&equivalent_expr.data_type(&self.schema)?,
)?,
}])?
.sort_properties
{
new_orderings.push(LexOrdering::new(ordering[1..].to_vec()));
break;
}
}
}
self.oeq_class.add_new_orderings(new_orderings);
Ok(())
}
pub fn with_reorder(mut self, sort_exprs: LexOrdering) -> Self {
self.oeq_class = OrderingEquivalenceClass::new(vec![sort_exprs]);
self
}
fn normalize_sort_exprs(&self, sort_exprs: LexOrderingRef) -> LexOrdering {
let sort_reqs = PhysicalSortRequirement::from_sort_exprs(sort_exprs.iter());
let normalized_sort_reqs = self.normalize_sort_requirements(&sort_reqs);
PhysicalSortRequirement::to_sort_exprs(normalized_sort_reqs)
}
fn normalize_sort_requirements(
&self,
sort_reqs: LexRequirementRef,
) -> LexRequirement {
let normalized_sort_reqs = self.eq_group.normalize_sort_requirements(sort_reqs);
let mut constant_exprs = vec![];
constant_exprs.extend(
self.constants
.iter()
.map(|const_expr| Arc::clone(const_expr.expr())),
);
let constants_normalized = self.eq_group.normalize_exprs(constant_exprs);
collapse_lex_req(
normalized_sort_reqs
.iter()
.filter(|&order| {
!physical_exprs_contains(&constants_normalized, &order.expr)
})
.cloned()
.collect(),
)
}
pub fn ordering_satisfy(&self, given: LexOrderingRef) -> bool {
let sort_requirements = PhysicalSortRequirement::from_sort_exprs(given.iter());
self.ordering_satisfy_requirement(&sort_requirements)
}
pub fn ordering_satisfy_requirement(&self, reqs: LexRequirementRef) -> bool {
let mut eq_properties = self.clone();
let normalized_reqs = eq_properties.normalize_sort_requirements(reqs);
for normalized_req in normalized_reqs {
if !eq_properties.ordering_satisfy_single(&normalized_req) {
return false;
}
eq_properties = eq_properties
.with_constants(std::iter::once(ConstExpr::from(normalized_req.expr)));
}
true
}
fn ordering_satisfy_single(&self, req: &PhysicalSortRequirement) -> bool {
let ExprProperties {
sort_properties, ..
} = self.get_expr_properties(Arc::clone(&req.expr));
match sort_properties {
SortProperties::Ordered(options) => {
let sort_expr = PhysicalSortExpr {
expr: Arc::clone(&req.expr),
options,
};
sort_expr.satisfy(req, self.schema())
}
SortProperties::Singleton => true,
SortProperties::Unordered => false,
}
}
pub fn requirements_compatible(
&self,
given: LexRequirementRef,
reference: LexRequirementRef,
) -> bool {
let normalized_given = self.normalize_sort_requirements(given);
let normalized_reference = self.normalize_sort_requirements(reference);
(normalized_reference.len() <= normalized_given.len())
&& normalized_reference
.into_iter()
.zip(normalized_given)
.all(|(reference, given)| given.compatible(&reference))
}
pub fn get_finer_ordering(
&self,
lhs: LexOrderingRef,
rhs: LexOrderingRef,
) -> Option<LexOrdering> {
let lhs = PhysicalSortRequirement::from_sort_exprs(lhs);
let rhs = PhysicalSortRequirement::from_sort_exprs(rhs);
let finer = self.get_finer_requirement(&lhs, &rhs);
finer.map(PhysicalSortRequirement::to_sort_exprs)
}
pub fn get_finer_requirement(
&self,
req1: LexRequirementRef,
req2: LexRequirementRef,
) -> Option<LexRequirement> {
let mut lhs = self.normalize_sort_requirements(req1);
let mut rhs = self.normalize_sort_requirements(req2);
lhs.inner
.iter_mut()
.zip(rhs.inner.iter_mut())
.all(|(lhs, rhs)| {
lhs.expr.eq(&rhs.expr)
&& match (lhs.options, rhs.options) {
(Some(lhs_opt), Some(rhs_opt)) => lhs_opt == rhs_opt,
(Some(options), None) => {
rhs.options = Some(options);
true
}
(None, Some(options)) => {
lhs.options = Some(options);
true
}
(None, None) => true,
}
})
.then_some(if lhs.len() >= rhs.len() { lhs } else { rhs })
}
pub fn substitute_ordering_component(
&self,
mapping: &ProjectionMapping,
sort_expr: LexOrderingRef,
) -> Result<Vec<LexOrdering>> {
let new_orderings = sort_expr
.iter()
.map(|sort_expr| {
let referring_exprs: Vec<_> = mapping
.iter()
.map(|(source, _target)| source)
.filter(|source| expr_refers(source, &sort_expr.expr))
.cloned()
.collect();
let mut res = LexOrdering::new(vec![sort_expr.clone()]);
for r_expr in referring_exprs {
if let Some(cast_expr) = r_expr.as_any().downcast_ref::<CastExpr>() {
let expr_type = sort_expr.expr.data_type(&self.schema)?;
if cast_expr.expr.eq(&sort_expr.expr)
&& cast_expr.is_bigger_cast(expr_type)
{
res.push(PhysicalSortExpr {
expr: Arc::clone(&r_expr),
options: sort_expr.options,
});
}
}
}
Ok(res)
})
.collect::<Result<Vec<_>>>()?;
let res = new_orderings
.into_iter()
.map(|ordering| ordering.inner)
.multi_cartesian_product()
.map(LexOrdering::new)
.collect::<Vec<_>>();
Ok(res)
}
pub fn substitute_oeq_class(&mut self, mapping: &ProjectionMapping) -> Result<()> {
let orderings = &self.oeq_class.orderings;
let new_order = orderings
.iter()
.map(|order| self.substitute_ordering_component(mapping, order.as_ref()))
.collect::<Result<Vec<_>>>()?;
let new_order = new_order.into_iter().flatten().collect();
self.oeq_class = OrderingEquivalenceClass::new(new_order);
Ok(())
}
pub fn project_expr(
&self,
expr: &Arc<dyn PhysicalExpr>,
projection_mapping: &ProjectionMapping,
) -> Option<Arc<dyn PhysicalExpr>> {
self.eq_group.project_expr(projection_mapping, expr)
}
fn construct_dependency_map(&self, mapping: &ProjectionMapping) -> DependencyMap {
let mut dependency_map = DependencyMap::new();
for ordering in self.normalized_oeq_class().iter() {
for (idx, sort_expr) in ordering.iter().enumerate() {
let target_sort_expr =
self.project_expr(&sort_expr.expr, mapping).map(|expr| {
PhysicalSortExpr {
expr,
options: sort_expr.options,
}
});
let is_projected = target_sort_expr.is_some();
if is_projected
|| mapping
.iter()
.any(|(source, _)| expr_refers(source, &sort_expr.expr))
{
let dependency = idx.checked_sub(1).map(|a| &ordering[a]);
dependency_map.insert(
sort_expr,
target_sort_expr.as_ref(),
dependency,
);
}
if !is_projected {
break;
}
}
}
dependency_map
}
fn normalized_mapping(&self, mapping: &ProjectionMapping) -> ProjectionMapping {
ProjectionMapping {
map: mapping
.iter()
.map(|(source, target)| {
let normalized_source =
self.eq_group.normalize_expr(Arc::clone(source));
(normalized_source, Arc::clone(target))
})
.collect(),
}
}
fn projected_orderings(&self, mapping: &ProjectionMapping) -> Vec<LexOrdering> {
let mapping = self.normalized_mapping(mapping);
let dependency_map = self.construct_dependency_map(&mapping);
let orderings = mapping.iter().flat_map(|(source, target)| {
referred_dependencies(&dependency_map, source)
.into_iter()
.filter_map(|relevant_deps| {
if let Ok(SortProperties::Ordered(options)) =
get_expr_properties(source, &relevant_deps, &self.schema)
.map(|prop| prop.sort_properties)
{
Some((options, relevant_deps))
} else {
None
}
})
.flat_map(|(options, relevant_deps)| {
let sort_expr = PhysicalSortExpr {
expr: Arc::clone(target),
options,
};
let mut dependency_orderings =
generate_dependency_orderings(&relevant_deps, &dependency_map);
for ordering in dependency_orderings.iter_mut() {
ordering.push(sort_expr.clone());
}
dependency_orderings
})
});
let projected_orderings = dependency_map.iter().flat_map(|(sort_expr, node)| {
let mut prefixes = construct_prefix_orderings(sort_expr, &dependency_map);
if prefixes.is_empty() {
prefixes = vec![LexOrdering::default()];
}
for ordering in prefixes.iter_mut() {
if let Some(target) = &node.target_sort_expr {
ordering.push(target.clone())
}
}
prefixes
});
orderings
.chain(projected_orderings)
.map(collapse_lex_ordering)
.collect()
}
fn projected_constants(&self, mapping: &ProjectionMapping) -> Vec<ConstExpr> {
let mut projected_constants = self
.constants
.iter()
.flat_map(|const_expr| {
const_expr.map(|expr| self.eq_group.project_expr(mapping, expr))
})
.collect::<Vec<_>>();
for (source, target) in mapping.iter() {
if self.is_expr_constant(source)
&& !const_exprs_contains(&projected_constants, target)
{
projected_constants
.push(ConstExpr::from(target).with_across_partitions(true));
}
}
projected_constants
}
pub fn project(
&self,
projection_mapping: &ProjectionMapping,
output_schema: SchemaRef,
) -> Self {
let projected_constants = self.projected_constants(projection_mapping);
let projected_eq_group = self.eq_group.project(projection_mapping);
let projected_orderings = self.projected_orderings(projection_mapping);
Self {
eq_group: projected_eq_group,
oeq_class: OrderingEquivalenceClass::new(projected_orderings),
constants: projected_constants,
schema: output_schema,
}
}
pub fn find_longest_permutation(
&self,
exprs: &[Arc<dyn PhysicalExpr>],
) -> (LexOrdering, Vec<usize>) {
let mut eq_properties = self.clone();
let mut result = vec![];
let mut search_indices = (0..exprs.len()).collect::<IndexSet<_>>();
for _idx in 0..exprs.len() {
let ordered_exprs = search_indices
.iter()
.flat_map(|&idx| {
let ExprProperties {
sort_properties, ..
} = eq_properties.get_expr_properties(Arc::clone(&exprs[idx]));
match sort_properties {
SortProperties::Ordered(options) => Some((
PhysicalSortExpr {
expr: Arc::clone(&exprs[idx]),
options,
},
idx,
)),
SortProperties::Singleton => {
let options = SortOptions::default();
Some((
PhysicalSortExpr {
expr: Arc::clone(&exprs[idx]),
options,
},
idx,
))
}
SortProperties::Unordered => None,
}
})
.collect::<Vec<_>>();
if ordered_exprs.is_empty() {
break;
}
for (PhysicalSortExpr { expr, .. }, idx) in &ordered_exprs {
eq_properties =
eq_properties.with_constants(std::iter::once(ConstExpr::from(expr)));
search_indices.shift_remove(idx);
}
result.extend(ordered_exprs);
}
let (left, right) = result.into_iter().unzip();
(LexOrdering::new(left), right)
}
pub fn is_expr_constant(&self, expr: &Arc<dyn PhysicalExpr>) -> bool {
let const_exprs = self
.constants
.iter()
.map(|const_expr| Arc::clone(const_expr.expr()));
let normalized_constants = self.eq_group.normalize_exprs(const_exprs);
let normalized_expr = self.eq_group.normalize_expr(Arc::clone(expr));
is_constant_recurse(&normalized_constants, &normalized_expr)
}
pub fn get_expr_properties(&self, expr: Arc<dyn PhysicalExpr>) -> ExprProperties {
ExprPropertiesNode::new_unknown(expr)
.transform_up(|expr| update_properties(expr, self))
.data()
.map(|node| node.data)
.unwrap_or(ExprProperties::new_unknown())
}
pub fn with_new_schema(self, schema: SchemaRef) -> Result<Self> {
let schemas_aligned = (self.schema.fields.len() == schema.fields.len())
&& self
.schema
.fields
.iter()
.zip(schema.fields.iter())
.all(|(lhs, rhs)| lhs.data_type().eq(rhs.data_type()));
if !schemas_aligned {
return plan_err!(
"Cannot rewrite old_schema:{:?} with new schema: {:?}",
self.schema,
schema
);
}
let new_constants = self
.constants
.into_iter()
.map(|const_expr| {
let across_partitions = const_expr.across_partitions();
let new_const_expr = with_new_schema(const_expr.owned_expr(), &schema)?;
Ok(ConstExpr::new(new_const_expr)
.with_across_partitions(across_partitions))
})
.collect::<Result<Vec<_>>>()?;
let mut new_orderings = vec![];
for ordering in self.oeq_class.orderings {
let new_ordering = ordering
.inner
.into_iter()
.map(|mut sort_expr| {
sort_expr.expr = with_new_schema(sort_expr.expr, &schema)?;
Ok(sort_expr)
})
.collect::<Result<_>>()?;
new_orderings.push(new_ordering);
}
let mut eq_classes = vec![];
for eq_class in self.eq_group.classes {
let new_eq_exprs = eq_class
.into_vec()
.into_iter()
.map(|expr| with_new_schema(expr, &schema))
.collect::<Result<_>>()?;
eq_classes.push(EquivalenceClass::new(new_eq_exprs));
}
let mut result = EquivalenceProperties::new(schema);
result.constants = new_constants;
result.add_new_orderings(new_orderings);
result.add_equivalence_group(EquivalenceGroup::new(eq_classes));
Ok(result)
}
}
impl Display for EquivalenceProperties {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
if self.eq_group.is_empty()
&& self.oeq_class.is_empty()
&& self.constants.is_empty()
{
return write!(f, "No properties");
}
if !self.oeq_class.is_empty() {
write!(f, "order: {}", self.oeq_class)?;
}
if !self.eq_group.is_empty() {
write!(f, ", eq: {}", self.eq_group)?;
}
if !self.constants.is_empty() {
write!(f, ", const: [{}]", ConstExpr::format_list(&self.constants))?;
}
Ok(())
}
}
fn update_properties(
mut node: ExprPropertiesNode,
eq_properties: &EquivalenceProperties,
) -> Result<Transformed<ExprPropertiesNode>> {
if !node.expr.children().is_empty() {
let children_props = node.children.iter().map(|c| c.data.clone()).collect_vec();
node.data = node.expr.get_properties(&children_props)?;
} else if node.expr.as_any().is::<Literal>() {
node.data = node.expr.get_properties(&[])?;
} else if node.expr.as_any().is::<Column>() {
node.data.range =
Interval::make_unbounded(&node.expr.data_type(eq_properties.schema())?)?
}
let normalized_expr = eq_properties
.eq_group
.normalize_expr(Arc::clone(&node.expr));
if eq_properties.is_expr_constant(&normalized_expr) {
node.data.sort_properties = SortProperties::Singleton;
} else if let Some(options) = eq_properties
.normalized_oeq_class()
.get_options(&normalized_expr)
{
node.data.sort_properties = SortProperties::Ordered(options);
}
Ok(Transformed::yes(node))
}
fn is_constant_recurse(
constants: &[Arc<dyn PhysicalExpr>],
expr: &Arc<dyn PhysicalExpr>,
) -> bool {
if physical_exprs_contains(constants, expr) || expr.as_any().is::<Literal>() {
return true;
}
let children = expr.children();
!children.is_empty() && children.iter().all(|c| is_constant_recurse(constants, c))
}
fn expr_refers(
referring_expr: &Arc<dyn PhysicalExpr>,
referred_expr: &Arc<dyn PhysicalExpr>,
) -> bool {
referring_expr.eq(referred_expr)
|| referring_expr
.children()
.iter()
.any(|child| expr_refers(child, referred_expr))
}
fn referred_dependencies(
dependency_map: &DependencyMap,
source: &Arc<dyn PhysicalExpr>,
) -> Vec<Dependencies> {
let mut expr_to_sort_exprs = IndexMap::<ExprWrapper, Dependencies>::new();
for sort_expr in dependency_map
.sort_exprs()
.filter(|sort_expr| expr_refers(source, &sort_expr.expr))
{
let key = ExprWrapper(Arc::clone(&sort_expr.expr));
expr_to_sort_exprs
.entry(key)
.or_default()
.insert(sort_expr.clone());
}
let dependencies = expr_to_sort_exprs
.into_values()
.map(Dependencies::into_inner)
.collect::<Vec<_>>();
dependencies
.iter()
.multi_cartesian_product()
.map(|referred_deps| {
Dependencies::new_from_iter(referred_deps.into_iter().cloned())
})
.collect()
}
fn construct_prefix_orderings(
relevant_sort_expr: &PhysicalSortExpr,
dependency_map: &DependencyMap,
) -> Vec<LexOrdering> {
let mut dep_enumerator = DependencyEnumerator::new();
dependency_map
.get(relevant_sort_expr)
.expect("no relevant sort expr found")
.dependencies
.iter()
.flat_map(|dep| dep_enumerator.construct_orderings(dep, dependency_map))
.collect()
}
fn generate_dependency_orderings(
dependencies: &Dependencies,
dependency_map: &DependencyMap,
) -> Vec<LexOrdering> {
let relevant_prefixes = dependencies
.iter()
.flat_map(|dep| {
let prefixes = construct_prefix_orderings(dep, dependency_map);
(!prefixes.is_empty()).then_some(prefixes)
})
.collect::<Vec<_>>();
if relevant_prefixes.is_empty() {
return vec![LexOrdering::default()];
}
relevant_prefixes
.into_iter()
.multi_cartesian_product()
.flat_map(|prefix_orderings| {
prefix_orderings
.iter()
.permutations(prefix_orderings.len())
.map(|prefixes| {
prefixes
.into_iter()
.flat_map(|ordering| ordering.inner.clone())
.collect()
})
.collect::<Vec<_>>()
})
.collect()
}
fn get_expr_properties(
expr: &Arc<dyn PhysicalExpr>,
dependencies: &Dependencies,
schema: &SchemaRef,
) -> Result<ExprProperties> {
if let Some(column_order) = dependencies.iter().find(|&order| expr.eq(&order.expr)) {
Ok(ExprProperties {
sort_properties: SortProperties::Ordered(column_order.options),
range: Interval::make_unbounded(&expr.data_type(schema)?)?,
})
} else if expr.as_any().downcast_ref::<Column>().is_some() {
Ok(ExprProperties {
sort_properties: SortProperties::Unordered,
range: Interval::make_unbounded(&expr.data_type(schema)?)?,
})
} else if let Some(literal) = expr.as_any().downcast_ref::<Literal>() {
Ok(ExprProperties {
sort_properties: SortProperties::Singleton,
range: Interval::try_new(literal.value().clone(), literal.value().clone())?,
})
} else {
let child_states = expr
.children()
.iter()
.map(|child| get_expr_properties(child, dependencies, schema))
.collect::<Result<Vec<_>>>()?;
expr.get_properties(&child_states)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct DependencyNode {
target_sort_expr: Option<PhysicalSortExpr>,
dependencies: Dependencies,
}
impl DependencyNode {
fn insert_dependency(&mut self, dependency: Option<&PhysicalSortExpr>) {
if let Some(dep) = dependency {
self.dependencies.insert(dep.clone());
}
}
}
impl Display for DependencyNode {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
if let Some(target) = &self.target_sort_expr {
write!(f, "(target: {}, ", target)?;
} else {
write!(f, "(")?;
}
write!(f, "dependencies: [{}])", self.dependencies)
}
}
#[derive(Debug)]
struct DependencyMap {
inner: IndexMap<PhysicalSortExpr, DependencyNode>,
}
impl DependencyMap {
fn new() -> Self {
Self {
inner: IndexMap::new(),
}
}
fn insert(
&mut self,
sort_expr: &PhysicalSortExpr,
target_sort_expr: Option<&PhysicalSortExpr>,
dependency: Option<&PhysicalSortExpr>,
) {
self.inner
.entry(sort_expr.clone())
.or_insert_with(|| DependencyNode {
target_sort_expr: target_sort_expr.cloned(),
dependencies: Dependencies::new(),
})
.insert_dependency(dependency)
}
fn iter(&self) -> impl Iterator<Item = (&PhysicalSortExpr, &DependencyNode)> {
self.inner.iter()
}
fn sort_exprs(&self) -> impl Iterator<Item = &PhysicalSortExpr> {
self.inner.keys()
}
fn get(&self, sort_expr: &PhysicalSortExpr) -> Option<&DependencyNode> {
self.inner.get(sort_expr)
}
}
impl Display for DependencyMap {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
writeln!(f, "DependencyMap: {{")?;
for (sort_expr, node) in self.inner.iter() {
writeln!(f, " {sort_expr} --> {node}")?;
}
writeln!(f, "}}")
}
}
#[derive(Debug, Default, Clone, PartialEq, Eq)]
struct Dependencies {
inner: IndexSet<PhysicalSortExpr>,
}
impl Display for Dependencies {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "[")?;
let mut iter = self.inner.iter();
if let Some(dep) = iter.next() {
write!(f, "{}", dep)?;
}
for dep in iter {
write!(f, ", {}", dep)?;
}
write!(f, "]")
}
}
impl Dependencies {
fn new() -> Self {
Self {
inner: IndexSet::new(),
}
}
fn new_from_iter(iter: impl IntoIterator<Item = PhysicalSortExpr>) -> Self {
Self {
inner: iter.into_iter().collect(),
}
}
fn insert(&mut self, sort_expr: PhysicalSortExpr) {
self.inner.insert(sort_expr);
}
fn iter(&self) -> impl Iterator<Item = &PhysicalSortExpr> + Clone {
self.inner.iter()
}
fn into_inner(self) -> IndexSet<PhysicalSortExpr> {
self.inner
}
fn is_empty(&self) -> bool {
self.inner.is_empty()
}
}
struct DependencyEnumerator<'a> {
seen: IndexMap<&'a PhysicalSortExpr, IndexSet<&'a PhysicalSortExpr>>,
}
impl<'a> DependencyEnumerator<'a> {
fn new() -> Self {
Self {
seen: IndexMap::new(),
}
}
fn insert(
&mut self,
target: &'a PhysicalSortExpr,
dep: &'a PhysicalSortExpr,
) -> bool {
self.seen.entry(target).or_default().insert(dep)
}
fn construct_orderings(
&mut self,
referred_sort_expr: &'a PhysicalSortExpr,
dependency_map: &'a DependencyMap,
) -> Vec<LexOrdering> {
let node = dependency_map
.get(referred_sort_expr)
.expect("`referred_sort_expr` should be inside `dependency_map`");
let target_sort_expr = node.target_sort_expr.as_ref().unwrap();
if node.dependencies.is_empty() {
return vec![LexOrdering::new(vec![target_sort_expr.clone()])];
};
node.dependencies
.iter()
.flat_map(|dep| {
let mut orderings = if self.insert(target_sort_expr, dep) {
self.construct_orderings(dep, dependency_map)
} else {
vec![]
};
for ordering in orderings.iter_mut() {
ordering.push(target_sort_expr.clone())
}
orderings
})
.collect()
}
}
pub fn join_equivalence_properties(
left: EquivalenceProperties,
right: EquivalenceProperties,
join_type: &JoinType,
join_schema: SchemaRef,
maintains_input_order: &[bool],
probe_side: Option<JoinSide>,
on: &[(PhysicalExprRef, PhysicalExprRef)],
) -> EquivalenceProperties {
let left_size = left.schema.fields.len();
let mut result = EquivalenceProperties::new(join_schema);
result.add_equivalence_group(left.eq_group().join(
right.eq_group(),
join_type,
left_size,
on,
));
let EquivalenceProperties {
constants: left_constants,
oeq_class: left_oeq_class,
..
} = left;
let EquivalenceProperties {
constants: right_constants,
oeq_class: mut right_oeq_class,
..
} = right;
match maintains_input_order {
[true, false] => {
if let (Some(JoinSide::Left), JoinType::Inner) = (probe_side, join_type) {
updated_right_ordering_equivalence_class(
&mut right_oeq_class,
join_type,
left_size,
);
let out_oeq_class = left_oeq_class.join_suffix(&right_oeq_class);
result.add_ordering_equivalence_class(out_oeq_class);
} else {
result.add_ordering_equivalence_class(left_oeq_class);
}
}
[false, true] => {
updated_right_ordering_equivalence_class(
&mut right_oeq_class,
join_type,
left_size,
);
if let (Some(JoinSide::Right), JoinType::Inner) = (probe_side, join_type) {
let out_oeq_class = right_oeq_class.join_suffix(&left_oeq_class);
result.add_ordering_equivalence_class(out_oeq_class);
} else {
result.add_ordering_equivalence_class(right_oeq_class);
}
}
[false, false] => {}
[true, true] => unreachable!("Cannot maintain ordering of both sides"),
_ => unreachable!("Join operators can not have more than two children"),
}
match join_type {
JoinType::LeftAnti | JoinType::LeftSemi => {
result = result.with_constants(left_constants);
}
JoinType::RightAnti | JoinType::RightSemi => {
result = result.with_constants(right_constants);
}
_ => {}
}
result
}
fn updated_right_ordering_equivalence_class(
right_oeq_class: &mut OrderingEquivalenceClass,
join_type: &JoinType,
left_size: usize,
) {
if matches!(
join_type,
JoinType::Inner | JoinType::Left | JoinType::Full | JoinType::Right
) {
right_oeq_class.add_offset(left_size);
}
}
#[derive(Debug, Clone)]
struct ExprWrapper(Arc<dyn PhysicalExpr>);
impl PartialEq<Self> for ExprWrapper {
fn eq(&self, other: &Self) -> bool {
self.0.eq(&other.0)
}
}
impl Eq for ExprWrapper {}
impl Hash for ExprWrapper {
fn hash<H: Hasher>(&self, state: &mut H) {
self.0.hash(state);
}
}
fn calculate_union_binary(
mut lhs: EquivalenceProperties,
mut rhs: EquivalenceProperties,
) -> Result<EquivalenceProperties> {
if !rhs.schema.eq(&lhs.schema) {
rhs = rhs.with_new_schema(Arc::clone(&lhs.schema))?;
}
let constants: Vec<_> = lhs
.constants()
.iter()
.filter(|const_expr| const_exprs_contains(rhs.constants(), const_expr.expr()))
.map(|const_expr| {
ConstExpr::new(Arc::clone(const_expr.expr())).with_across_partitions(false)
})
.collect();
for c in &constants {
lhs = lhs.remove_constant(c);
rhs = rhs.remove_constant(c);
}
let mut orderings = UnionEquivalentOrderingBuilder::new();
orderings.add_satisfied_orderings(
lhs.normalized_oeq_class().orderings,
lhs.constants(),
&rhs,
);
orderings.add_satisfied_orderings(
rhs.normalized_oeq_class().orderings,
rhs.constants(),
&lhs,
);
let orderings = orderings.build();
let mut eq_properties =
EquivalenceProperties::new(lhs.schema).with_constants(constants);
eq_properties.add_new_orderings(orderings);
Ok(eq_properties)
}
pub fn calculate_union(
eqps: Vec<EquivalenceProperties>,
schema: SchemaRef,
) -> Result<EquivalenceProperties> {
let mut iter = eqps.into_iter();
let Some(mut acc) = iter.next() else {
return internal_err!(
"Cannot calculate EquivalenceProperties for a union with no inputs"
);
};
if !acc.schema.eq(&schema) {
acc = acc.with_new_schema(schema)?;
}
for props in iter {
acc = calculate_union_binary(acc, props)?;
}
Ok(acc)
}
#[derive(Debug)]
enum AddedOrdering {
Yes,
No(LexOrdering),
}
#[derive(Debug)]
struct UnionEquivalentOrderingBuilder {
orderings: Vec<LexOrdering>,
}
impl UnionEquivalentOrderingBuilder {
fn new() -> Self {
Self { orderings: vec![] }
}
fn add_satisfied_orderings(
&mut self,
orderings: impl IntoIterator<Item = LexOrdering>,
constants: &[ConstExpr],
properties: &EquivalenceProperties,
) {
for mut ordering in orderings.into_iter() {
loop {
match self.try_add_ordering(ordering, constants, properties) {
AddedOrdering::Yes => break,
AddedOrdering::No(o) => {
ordering = o;
ordering.pop();
}
}
}
}
}
fn try_add_ordering(
&mut self,
ordering: LexOrdering,
constants: &[ConstExpr],
properties: &EquivalenceProperties,
) -> AddedOrdering {
if ordering.is_empty() {
AddedOrdering::Yes
} else if constants.is_empty() && properties.ordering_satisfy(ordering.as_ref()) {
self.orderings.push(ordering);
AddedOrdering::Yes
} else {
if self.try_find_augmented_ordering(&ordering, constants, properties) {
AddedOrdering::Yes
} else {
AddedOrdering::No(ordering)
}
}
}
fn try_find_augmented_ordering(
&mut self,
ordering: &LexOrdering,
constants: &[ConstExpr],
properties: &EquivalenceProperties,
) -> bool {
if constants.is_empty() {
return false;
}
let start_num_orderings = self.orderings.len();
for existing_ordering in &properties.oeq_class.orderings {
if let Some(augmented_ordering) = self.augment_ordering(
ordering,
constants,
existing_ordering,
&properties.constants,
) {
if !augmented_ordering.is_empty() {
assert!(properties.ordering_satisfy(augmented_ordering.as_ref()));
self.orderings.push(augmented_ordering);
}
}
}
self.orderings.len() > start_num_orderings
}
fn augment_ordering(
&mut self,
ordering: &LexOrdering,
constants: &[ConstExpr],
existing_ordering: &LexOrdering,
existing_constants: &[ConstExpr],
) -> Option<LexOrdering> {
let mut augmented_ordering = LexOrdering::default();
let mut sort_expr_iter = ordering.inner.iter().peekable();
let mut existing_sort_expr_iter = existing_ordering.inner.iter().peekable();
while sort_expr_iter.peek().is_some() || existing_sort_expr_iter.peek().is_some()
{
if let Some(expr) =
advance_if_match(&mut sort_expr_iter, &mut existing_sort_expr_iter)
{
augmented_ordering.push(expr);
} else if let Some(expr) =
advance_if_matches_constant(&mut sort_expr_iter, existing_constants)
{
augmented_ordering.push(expr);
} else if let Some(expr) =
advance_if_matches_constant(&mut existing_sort_expr_iter, constants)
{
augmented_ordering.push(expr);
} else {
break;
}
}
Some(augmented_ordering)
}
fn build(self) -> Vec<LexOrdering> {
self.orderings
}
}
fn advance_if_match(
iter1: &mut Peekable<Iter<PhysicalSortExpr>>,
iter2: &mut Peekable<Iter<PhysicalSortExpr>>,
) -> Option<PhysicalSortExpr> {
if matches!((iter1.peek(), iter2.peek()), (Some(expr1), Some(expr2)) if expr1.eq(expr2))
{
iter1.next().unwrap();
iter2.next().cloned()
} else {
None
}
}
fn advance_if_matches_constant(
iter: &mut Peekable<Iter<PhysicalSortExpr>>,
constants: &[ConstExpr],
) -> Option<PhysicalSortExpr> {
let expr = iter.peek()?;
let const_expr = constants.iter().find(|c| c.eq_expr(expr))?;
let found_expr = PhysicalSortExpr::new(Arc::clone(const_expr.expr()), expr.options);
iter.next();
Some(found_expr)
}
#[cfg(test)]
mod tests {
use std::ops::Not;
use super::*;
use crate::equivalence::add_offset_to_expr;
use crate::equivalence::tests::{
convert_to_orderings, convert_to_sort_exprs, convert_to_sort_reqs,
create_test_params, create_test_schema, output_schema,
};
use crate::expressions::{col, BinaryExpr, Column};
use arrow::datatypes::{DataType, Field, Schema};
use arrow_schema::{Fields, TimeUnit};
use datafusion_expr::Operator;
#[test]
fn project_equivalence_properties_test() -> Result<()> {
let input_schema = Arc::new(Schema::new(vec![
Field::new("a", DataType::Int64, true),
Field::new("b", DataType::Int64, true),
Field::new("c", DataType::Int64, true),
]));
let input_properties = EquivalenceProperties::new(Arc::clone(&input_schema));
let col_a = col("a", &input_schema)?;
let proj_exprs = vec![
(Arc::clone(&col_a), "a1".to_string()),
(Arc::clone(&col_a), "a2".to_string()),
(Arc::clone(&col_a), "a3".to_string()),
(Arc::clone(&col_a), "a4".to_string()),
];
let projection_mapping = ProjectionMapping::try_new(&proj_exprs, &input_schema)?;
let out_schema = output_schema(&projection_mapping, &input_schema)?;
let proj_exprs = vec![
(Arc::clone(&col_a), "a1".to_string()),
(Arc::clone(&col_a), "a2".to_string()),
(Arc::clone(&col_a), "a3".to_string()),
(Arc::clone(&col_a), "a4".to_string()),
];
let projection_mapping = ProjectionMapping::try_new(&proj_exprs, &input_schema)?;
let col_a1 = &col("a1", &out_schema)?;
let col_a2 = &col("a2", &out_schema)?;
let col_a3 = &col("a3", &out_schema)?;
let col_a4 = &col("a4", &out_schema)?;
let out_properties = input_properties.project(&projection_mapping, out_schema);
assert_eq!(out_properties.eq_group().len(), 1);
let eq_class = &out_properties.eq_group().classes[0];
assert_eq!(eq_class.len(), 4);
assert!(eq_class.contains(col_a1));
assert!(eq_class.contains(col_a2));
assert!(eq_class.contains(col_a3));
assert!(eq_class.contains(col_a4));
Ok(())
}
#[test]
fn project_equivalence_properties_test_multi() -> Result<()> {
let input_schema = Arc::new(Schema::new(vec![
Field::new("a", DataType::Int64, true),
Field::new("b", DataType::Int64, true),
Field::new("c", DataType::Int64, true),
Field::new("d", DataType::Int64, true),
]));
let mut input_properties = EquivalenceProperties::new(Arc::clone(&input_schema));
input_properties.add_new_ordering(LexOrdering::new(vec![
parse_sort_expr("a", &input_schema),
parse_sort_expr("b", &input_schema),
parse_sort_expr("c", &input_schema),
parse_sort_expr("d", &input_schema),
]));
input_properties.add_new_ordering(LexOrdering::new(vec![
parse_sort_expr("a", &input_schema),
parse_sort_expr("c", &input_schema),
parse_sort_expr("b", &input_schema), parse_sort_expr("d", &input_schema),
]));
let proj_exprs = vec![
(col("a", &input_schema)?, "a".to_string()),
(col("b", &input_schema)?, "b".to_string()),
(col("c", &input_schema)?, "c".to_string()),
(col("d", &input_schema)?, "d".to_string()),
];
let projection_mapping = ProjectionMapping::try_new(&proj_exprs, &input_schema)?;
let out_properties = input_properties.project(&projection_mapping, input_schema);
assert_eq!(
out_properties.to_string(),
"order: [[a@0 ASC, c@2 ASC, b@1 ASC, d@3 ASC], [a@0 ASC, b@1 ASC, c@2 ASC, d@3 ASC]]"
);
Ok(())
}
#[test]
fn test_join_equivalence_properties() -> 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 offset = schema.fields.len();
let col_a2 = &add_offset_to_expr(Arc::clone(col_a), offset);
let col_b2 = &add_offset_to_expr(Arc::clone(col_b), offset);
let option_asc = SortOptions {
descending: false,
nulls_first: false,
};
let test_cases = vec![
(
vec![vec![(col_a, option_asc)], vec![(col_b, option_asc)]],
vec![vec![(col_a, option_asc)], vec![(col_b, option_asc)]],
vec![
vec![(col_a, option_asc), (col_a2, option_asc)],
vec![(col_a, option_asc), (col_b2, option_asc)],
vec![(col_b, option_asc), (col_a2, option_asc)],
vec![(col_b, option_asc), (col_b2, option_asc)],
],
),
(
vec![
vec![(col_a, option_asc)],
vec![(col_b, option_asc)],
vec![(col_c, option_asc)],
],
vec![vec![(col_a, option_asc)], vec![(col_b, option_asc)]],
vec![
vec![(col_a, option_asc), (col_a2, option_asc)],
vec![(col_a, option_asc), (col_b2, option_asc)],
vec![(col_b, option_asc), (col_a2, option_asc)],
vec![(col_b, option_asc), (col_b2, option_asc)],
vec![(col_c, option_asc), (col_a2, option_asc)],
vec![(col_c, option_asc), (col_b2, option_asc)],
],
),
];
for (left_orderings, right_orderings, expected) in test_cases {
let mut left_eq_properties = EquivalenceProperties::new(Arc::clone(&schema));
let mut right_eq_properties = EquivalenceProperties::new(Arc::clone(&schema));
let left_orderings = convert_to_orderings(&left_orderings);
let right_orderings = convert_to_orderings(&right_orderings);
let expected = convert_to_orderings(&expected);
left_eq_properties.add_new_orderings(left_orderings);
right_eq_properties.add_new_orderings(right_orderings);
let join_eq = join_equivalence_properties(
left_eq_properties,
right_eq_properties,
&JoinType::Inner,
Arc::new(Schema::empty()),
&[true, false],
Some(JoinSide::Left),
&[],
);
let orderings = &join_eq.oeq_class.orderings;
let err_msg = format!("expected: {:?}, actual:{:?}", expected, orderings);
assert_eq!(
join_eq.oeq_class.orderings.len(),
expected.len(),
"{}",
err_msg
);
for ordering in orderings {
assert!(
expected.contains(ordering),
"{}, ordering: {:?}",
err_msg,
ordering
);
}
}
Ok(())
}
#[test]
fn test_expr_consists_of_constants() -> Result<()> {
let schema = Arc::new(Schema::new(vec![
Field::new("a", DataType::Int32, true),
Field::new("b", DataType::Int32, true),
Field::new("c", DataType::Int32, true),
Field::new("d", DataType::Int32, true),
Field::new("ts", DataType::Timestamp(TimeUnit::Nanosecond, None), true),
]));
let col_a = col("a", &schema)?;
let col_b = col("b", &schema)?;
let col_d = col("d", &schema)?;
let b_plus_d = Arc::new(BinaryExpr::new(
Arc::clone(&col_b),
Operator::Plus,
Arc::clone(&col_d),
)) as Arc<dyn PhysicalExpr>;
let constants = vec![Arc::clone(&col_a), Arc::clone(&col_b)];
let expr = Arc::clone(&b_plus_d);
assert!(!is_constant_recurse(&constants, &expr));
let constants = vec![Arc::clone(&col_a), Arc::clone(&col_b), Arc::clone(&col_d)];
let expr = Arc::clone(&b_plus_d);
assert!(is_constant_recurse(&constants, &expr));
Ok(())
}
#[test]
fn test_get_updated_right_ordering_equivalence_properties() -> Result<()> {
let join_type = JoinType::Inner;
let child_fields: Fields = ["x", "y", "z", "w"]
.into_iter()
.map(|name| Field::new(name, DataType::Int32, true))
.collect();
let child_schema = Schema::new(child_fields);
let col_x = &col("x", &child_schema)?;
let col_y = &col("y", &child_schema)?;
let col_z = &col("z", &child_schema)?;
let col_w = &col("w", &child_schema)?;
let option_asc = SortOptions {
descending: false,
nulls_first: false,
};
let orderings = vec![
vec![(col_x, option_asc), (col_y, option_asc)],
vec![(col_z, option_asc), (col_w, option_asc)],
];
let orderings = convert_to_orderings(&orderings);
let mut right_oeq_class = OrderingEquivalenceClass::new(orderings);
let left_columns_len = 4;
let fields: Fields = ["a", "b", "c", "d", "x", "y", "z", "w"]
.into_iter()
.map(|name| Field::new(name, DataType::Int32, true))
.collect();
let schema = Schema::new(fields);
let col_a = &col("a", &schema)?;
let col_d = &col("d", &schema)?;
let col_x = &col("x", &schema)?;
let col_y = &col("y", &schema)?;
let col_z = &col("z", &schema)?;
let col_w = &col("w", &schema)?;
let mut join_eq_properties = EquivalenceProperties::new(Arc::new(schema));
join_eq_properties.add_equal_conditions(col_a, col_x)?;
join_eq_properties.add_equal_conditions(col_d, col_w)?;
updated_right_ordering_equivalence_class(
&mut right_oeq_class,
&join_type,
left_columns_len,
);
join_eq_properties.add_ordering_equivalence_class(right_oeq_class);
let result = join_eq_properties.oeq_class().clone();
let orderings = vec![
vec![(col_x, option_asc), (col_y, option_asc)],
vec![(col_z, option_asc), (col_w, option_asc)],
];
let orderings = convert_to_orderings(&orderings);
let expected = OrderingEquivalenceClass::new(orderings);
assert_eq!(result, expected);
Ok(())
}
#[test]
fn test_normalize_ordering_equivalence_classes() -> Result<()> {
let sort_options = SortOptions::default();
let schema = Schema::new(vec![
Field::new("a", DataType::Int32, true),
Field::new("b", DataType::Int32, true),
Field::new("c", DataType::Int32, true),
]);
let col_a_expr = col("a", &schema)?;
let col_b_expr = col("b", &schema)?;
let col_c_expr = col("c", &schema)?;
let mut eq_properties = EquivalenceProperties::new(Arc::new(schema.clone()));
eq_properties.add_equal_conditions(&col_a_expr, &col_c_expr)?;
let others = vec![
LexOrdering::new(vec![PhysicalSortExpr {
expr: Arc::clone(&col_b_expr),
options: sort_options,
}]),
LexOrdering::new(vec![PhysicalSortExpr {
expr: Arc::clone(&col_c_expr),
options: sort_options,
}]),
];
eq_properties.add_new_orderings(others);
let mut expected_eqs = EquivalenceProperties::new(Arc::new(schema));
expected_eqs.add_new_orderings([
LexOrdering::new(vec![PhysicalSortExpr {
expr: Arc::clone(&col_b_expr),
options: sort_options,
}]),
LexOrdering::new(vec![PhysicalSortExpr {
expr: Arc::clone(&col_c_expr),
options: sort_options,
}]),
]);
let oeq_class = eq_properties.oeq_class().clone();
let expected = expected_eqs.oeq_class();
assert!(oeq_class.eq(expected));
Ok(())
}
#[test]
fn test_get_indices_of_matching_sort_exprs_with_order_eq() -> Result<()> {
let sort_options = SortOptions::default();
let sort_options_not = SortOptions::default().not();
let schema = Schema::new(vec![
Field::new("a", DataType::Int32, true),
Field::new("b", DataType::Int32, true),
]);
let col_a = &col("a", &schema)?;
let col_b = &col("b", &schema)?;
let required_columns = [Arc::clone(col_b), Arc::clone(col_a)];
let mut eq_properties = EquivalenceProperties::new(Arc::new(schema));
eq_properties.add_new_orderings([LexOrdering::new(vec![
PhysicalSortExpr {
expr: Arc::new(Column::new("b", 1)),
options: sort_options_not,
},
PhysicalSortExpr {
expr: Arc::new(Column::new("a", 0)),
options: sort_options,
},
])]);
let (result, idxs) = eq_properties.find_longest_permutation(&required_columns);
assert_eq!(idxs, vec![0, 1]);
assert_eq!(
result,
LexOrdering::new(vec![
PhysicalSortExpr {
expr: Arc::clone(col_b),
options: sort_options_not
},
PhysicalSortExpr {
expr: Arc::clone(col_a),
options: sort_options
}
])
);
let schema = Schema::new(vec![
Field::new("a", DataType::Int32, true),
Field::new("b", DataType::Int32, true),
Field::new("c", DataType::Int32, true),
]);
let col_a = &col("a", &schema)?;
let col_b = &col("b", &schema)?;
let required_columns = [Arc::clone(col_b), Arc::clone(col_a)];
let mut eq_properties = EquivalenceProperties::new(Arc::new(schema));
eq_properties.add_new_orderings([
LexOrdering::new(vec![PhysicalSortExpr {
expr: Arc::new(Column::new("c", 2)),
options: sort_options,
}]),
LexOrdering::new(vec![
PhysicalSortExpr {
expr: Arc::new(Column::new("b", 1)),
options: sort_options_not,
},
PhysicalSortExpr {
expr: Arc::new(Column::new("a", 0)),
options: sort_options,
},
]),
]);
let (result, idxs) = eq_properties.find_longest_permutation(&required_columns);
assert_eq!(idxs, vec![0, 1]);
assert_eq!(
result,
LexOrdering::new(vec![
PhysicalSortExpr {
expr: Arc::clone(col_b),
options: sort_options_not
},
PhysicalSortExpr {
expr: Arc::clone(col_a),
options: sort_options
}
])
);
let required_columns = [
Arc::new(Column::new("b", 1)) as _,
Arc::new(Column::new("a", 0)) as _,
];
let schema = Schema::new(vec![
Field::new("a", DataType::Int32, true),
Field::new("b", DataType::Int32, true),
Field::new("c", DataType::Int32, true),
]);
let mut eq_properties = EquivalenceProperties::new(Arc::new(schema));
eq_properties.add_new_orderings([LexOrdering::new(vec![
PhysicalSortExpr {
expr: Arc::new(Column::new("b", 1)),
options: sort_options_not,
},
PhysicalSortExpr {
expr: Arc::new(Column::new("c", 2)),
options: sort_options,
},
PhysicalSortExpr {
expr: Arc::new(Column::new("a", 0)),
options: sort_options,
},
])]);
let (_, idxs) = eq_properties.find_longest_permutation(&required_columns);
assert_eq!(idxs, vec![0]);
Ok(())
}
#[test]
fn test_update_properties() -> Result<()> {
let schema = Schema::new(vec![
Field::new("a", DataType::Int32, true),
Field::new("b", DataType::Int32, true),
Field::new("c", DataType::Int32, true),
Field::new("d", DataType::Int32, true),
]);
let mut eq_properties = EquivalenceProperties::new(Arc::new(schema.clone()));
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 option_asc = SortOptions {
descending: false,
nulls_first: false,
};
eq_properties.add_equal_conditions(col_b, col_a)?;
eq_properties.add_new_orderings(vec![
LexOrdering::new(vec![PhysicalSortExpr {
expr: Arc::clone(col_b),
options: option_asc,
}]),
LexOrdering::new(vec![PhysicalSortExpr {
expr: Arc::clone(col_d),
options: option_asc,
}]),
]);
let test_cases = vec![
(
Arc::new(BinaryExpr::new(
Arc::clone(col_d),
Operator::Plus,
Arc::clone(col_b),
)) as Arc<dyn PhysicalExpr>,
SortProperties::Ordered(option_asc),
),
(Arc::clone(col_b), SortProperties::Ordered(option_asc)),
(Arc::clone(col_a), SortProperties::Ordered(option_asc)),
(
Arc::new(BinaryExpr::new(
Arc::clone(col_a),
Operator::Plus,
Arc::clone(col_c),
)),
SortProperties::Unordered,
),
];
for (expr, expected) in test_cases {
let leading_orderings = eq_properties
.oeq_class()
.iter()
.flat_map(|ordering| ordering.inner.first().cloned())
.collect::<Vec<_>>();
let expr_props = eq_properties.get_expr_properties(Arc::clone(&expr));
let err_msg = format!(
"expr:{:?}, expected: {:?}, actual: {:?}, leading_orderings: {leading_orderings:?}",
expr, expected, expr_props.sort_properties
);
assert_eq!(expr_props.sort_properties, expected, "{}", err_msg);
}
Ok(())
}
#[test]
fn test_find_longest_permutation() -> Result<()> {
let (test_schema, mut eq_properties) = create_test_params()?;
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 col_h = &col("h", &test_schema)?;
let a_plus_d = Arc::new(BinaryExpr::new(
Arc::clone(col_a),
Operator::Plus,
Arc::clone(col_d),
)) as Arc<dyn PhysicalExpr>;
let option_asc = SortOptions {
descending: false,
nulls_first: false,
};
let option_desc = SortOptions {
descending: true,
nulls_first: true,
};
eq_properties.add_new_orderings([LexOrdering::new(vec![
PhysicalSortExpr {
expr: Arc::clone(col_d),
options: option_asc,
},
PhysicalSortExpr {
expr: Arc::clone(col_h),
options: option_desc,
},
])]);
let test_cases = vec![
(vec![col_a], vec![(col_a, option_asc)]),
(vec![col_c], vec![(col_c, option_asc)]),
(
vec![col_d, col_e, col_b],
vec![
(col_d, option_asc),
(col_e, option_desc),
(col_b, option_asc),
],
),
(vec![col_b], vec![]),
(vec![col_d], vec![(col_d, option_asc)]),
(vec![&a_plus_d], vec![(&a_plus_d, option_asc)]),
(
vec![col_b, col_d],
vec![(col_d, option_asc), (col_b, option_asc)],
),
(
vec![col_c, col_e],
vec![(col_c, option_asc), (col_e, option_desc)],
),
(
vec![col_d, col_h, col_e, col_f, col_b],
vec![
(col_d, option_asc),
(col_e, option_desc),
(col_h, option_desc),
(col_f, option_asc),
(col_b, option_asc),
],
),
(
vec![col_e, col_d, col_h, col_f, col_b],
vec![
(col_e, option_desc),
(col_d, option_asc),
(col_h, option_desc),
(col_f, option_asc),
(col_b, option_asc),
],
),
(
vec![col_e, col_d, col_b, col_h, col_f],
vec![
(col_e, option_desc),
(col_d, option_asc),
(col_b, option_asc),
(col_h, option_desc),
(col_f, option_asc),
],
),
];
for (exprs, expected) in test_cases {
let exprs = exprs.into_iter().cloned().collect::<Vec<_>>();
let expected = convert_to_sort_exprs(&expected);
let (actual, _) = eq_properties.find_longest_permutation(&exprs);
assert_eq!(actual, expected);
}
Ok(())
}
#[test]
fn test_find_longest_permutation2() -> Result<()> {
let (test_schema, mut eq_properties) = create_test_params()?;
let col_h = &col("h", &test_schema)?;
eq_properties = eq_properties.with_constants(vec![ConstExpr::from(col_h)]);
let test_cases = vec![
(vec![col_h], vec![(col_h, SortOptions::default())]),
];
for (exprs, expected) in test_cases {
let exprs = exprs.into_iter().cloned().collect::<Vec<_>>();
let expected = convert_to_sort_exprs(&expected);
let (actual, _) = eq_properties.find_longest_permutation(&exprs);
assert_eq!(actual, expected);
}
Ok(())
}
#[test]
fn test_get_finer() -> 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 eq_properties = EquivalenceProperties::new(schema);
let option_asc = SortOptions {
descending: false,
nulls_first: false,
};
let option_desc = SortOptions {
descending: true,
nulls_first: true,
};
let tests_cases = vec![
(
vec![(col_a, Some(option_asc))],
vec![(col_a, None), (col_b, Some(option_asc))],
Some(vec![(col_a, Some(option_asc)), (col_b, Some(option_asc))]),
),
(
vec![
(col_a, Some(option_asc)),
(col_b, Some(option_asc)),
(col_c, Some(option_asc)),
],
vec![(col_a, Some(option_asc)), (col_b, Some(option_asc))],
Some(vec![
(col_a, Some(option_asc)),
(col_b, Some(option_asc)),
(col_c, Some(option_asc)),
]),
),
(
vec![(col_a, Some(option_asc)), (col_b, Some(option_asc))],
vec![(col_a, Some(option_asc)), (col_b, Some(option_desc))],
None,
),
];
for (lhs, rhs, expected) in tests_cases {
let lhs = convert_to_sort_reqs(&lhs);
let rhs = convert_to_sort_reqs(&rhs);
let expected = expected.map(|expected| convert_to_sort_reqs(&expected));
let finer = eq_properties.get_finer_requirement(&lhs, &rhs);
assert_eq!(finer, expected)
}
Ok(())
}
#[test]
fn test_normalize_sort_reqs() -> Result<()> {
let (test_schema, eq_properties) = create_test_params()?;
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 option_asc = SortOptions {
descending: false,
nulls_first: false,
};
let option_desc = SortOptions {
descending: true,
nulls_first: true,
};
let requirements = vec![
(
vec![(col_a, Some(option_asc))],
vec![(col_a, Some(option_asc))],
),
(
vec![(col_a, Some(option_desc))],
vec![(col_a, Some(option_desc))],
),
(vec![(col_a, None)], vec![(col_a, None)]),
(
vec![(col_c, Some(option_asc))],
vec![(col_a, Some(option_asc))],
),
(vec![(col_c, None)], vec![(col_a, None)]),
(
vec![(col_d, Some(option_asc)), (col_b, Some(option_asc))],
vec![(col_d, Some(option_asc)), (col_b, Some(option_asc))],
),
(
vec![(col_d, None), (col_b, None)],
vec![(col_d, None), (col_b, None)],
),
(
vec![(col_e, Some(option_desc)), (col_f, Some(option_asc))],
vec![(col_e, Some(option_desc)), (col_f, Some(option_asc))],
),
(
vec![(col_e, Some(option_desc)), (col_f, None)],
vec![(col_e, Some(option_desc)), (col_f, None)],
),
(
vec![(col_e, None), (col_f, None)],
vec![(col_e, None), (col_f, None)],
),
];
for (reqs, expected_normalized) in requirements.into_iter() {
let req = convert_to_sort_reqs(&reqs);
let expected_normalized = convert_to_sort_reqs(&expected_normalized);
assert_eq!(
eq_properties.normalize_sort_requirements(&req),
expected_normalized
);
}
Ok(())
}
#[test]
fn test_schema_normalize_sort_requirement_with_equivalence() -> Result<()> {
let option1 = SortOptions {
descending: false,
nulls_first: false,
};
let (test_schema, eq_properties) = create_test_params()?;
let col_a = &col("a", &test_schema)?;
let col_c = &col("c", &test_schema)?;
let col_d = &col("d", &test_schema)?;
let test_cases = vec![
(vec![(col_a, Some(option1))], vec![(col_a, Some(option1))]),
(vec![(col_c, Some(option1))], vec![(col_a, Some(option1))]),
(vec![(col_c, None)], vec![(col_a, None)]),
(vec![(col_d, Some(option1))], vec![(col_d, Some(option1))]),
];
for (reqs, expected) in test_cases.into_iter() {
let reqs = convert_to_sort_reqs(&reqs);
let expected = convert_to_sort_reqs(&expected);
let normalized = eq_properties.normalize_sort_requirements(&reqs);
assert!(
expected.eq(&normalized),
"error in test: reqs: {reqs:?}, expected: {expected:?}, normalized: {normalized:?}"
);
}
Ok(())
}
#[test]
fn test_eliminate_redundant_monotonic_sorts() -> Result<()> {
let schema = Arc::new(Schema::new(vec![
Field::new("a", DataType::Date32, true),
Field::new("b", DataType::Utf8, true),
Field::new("c", DataType::Timestamp(TimeUnit::Nanosecond, None), true),
]));
let base_properties = EquivalenceProperties::new(Arc::clone(&schema))
.with_reorder(LexOrdering::new(
["a", "b", "c"]
.into_iter()
.map(|c| {
col(c, schema.as_ref()).map(|expr| PhysicalSortExpr {
expr,
options: SortOptions {
descending: false,
nulls_first: true,
},
})
})
.collect::<Result<Vec<_>>>()?,
));
struct TestCase {
name: &'static str,
constants: Vec<Arc<dyn PhysicalExpr>>,
equal_conditions: Vec<[Arc<dyn PhysicalExpr>; 2]>,
sort_columns: &'static [&'static str],
should_satisfy_ordering: bool,
}
let col_a = col("a", schema.as_ref())?;
let col_b = col("b", schema.as_ref())?;
let col_c = col("c", schema.as_ref())?;
let cast_c = Arc::new(CastExpr::new(col_c, DataType::Date32, None));
let cases = vec![
TestCase {
name: "(a, b, c) -> (c)",
constants: vec![Arc::clone(&col_b)],
equal_conditions: vec![[
Arc::clone(&cast_c) as Arc<dyn PhysicalExpr>,
Arc::clone(&col_a),
]],
sort_columns: &["c"],
should_satisfy_ordering: true,
},
TestCase {
name: "(a, b, c) -> (c)",
constants: vec![col_b],
equal_conditions: vec![[
Arc::clone(&col_a),
Arc::clone(&cast_c) as Arc<dyn PhysicalExpr>,
]],
sort_columns: &["c"],
should_satisfy_ordering: true,
},
TestCase {
name: "not ordered because (b) is not constant",
constants: vec![],
equal_conditions: vec![[
Arc::clone(&cast_c) as Arc<dyn PhysicalExpr>,
Arc::clone(&col_a),
]],
sort_columns: &["c"],
should_satisfy_ordering: false,
},
];
for case in cases {
for properties in [
{
let mut properties = base_properties.clone();
for [left, right] in &case.equal_conditions {
properties.add_equal_conditions(left, right)?
}
properties.with_constants(
case.constants.iter().cloned().map(ConstExpr::from),
)
},
{
let mut properties = base_properties.clone().with_constants(
case.constants.iter().cloned().map(ConstExpr::from),
);
for [left, right] in &case.equal_conditions {
properties.add_equal_conditions(left, right)?
}
properties
},
] {
let sort = case
.sort_columns
.iter()
.map(|&name| {
col(name, &schema).map(|col| PhysicalSortExpr {
expr: col,
options: SortOptions::default(),
})
})
.collect::<Result<LexOrdering>>()?;
assert_eq!(
properties.ordering_satisfy(sort.as_ref()),
case.should_satisfy_ordering,
"failed test '{}'",
case.name
);
}
}
Ok(())
}
fn append_fields(schema: &SchemaRef, text: &str) -> SchemaRef {
Arc::new(Schema::new(
schema
.fields()
.iter()
.map(|field| {
Field::new(
format!("{}{}", field.name(), text),
field.data_type().clone(),
field.is_nullable(),
)
})
.collect::<Vec<_>>(),
))
}
#[test]
fn test_union_equivalence_properties_multi_children_1() {
let schema = create_test_schema().unwrap();
let schema2 = append_fields(&schema, "1");
let schema3 = append_fields(&schema, "2");
UnionEquivalenceTest::new(&schema)
.with_child_sort(vec![vec!["a", "b", "c"]], &schema)
.with_child_sort(vec![vec!["a1", "b1", "c1"]], &schema2)
.with_child_sort(vec![vec!["a2", "b2"]], &schema3)
.with_expected_sort(vec![vec!["a", "b"]])
.run()
}
#[test]
fn test_union_equivalence_properties_multi_children_2() {
let schema = create_test_schema().unwrap();
let schema2 = append_fields(&schema, "1");
let schema3 = append_fields(&schema, "2");
UnionEquivalenceTest::new(&schema)
.with_child_sort(vec![vec!["a", "b", "c"]], &schema)
.with_child_sort(vec![vec!["a1", "b1", "c1"]], &schema2)
.with_child_sort(vec![vec!["a2", "b2", "c2"]], &schema3)
.with_expected_sort(vec![vec!["a", "b", "c"]])
.run()
}
#[test]
fn test_union_equivalence_properties_multi_children_3() {
let schema = create_test_schema().unwrap();
let schema2 = append_fields(&schema, "1");
let schema3 = append_fields(&schema, "2");
UnionEquivalenceTest::new(&schema)
.with_child_sort(vec![vec!["a", "b"]], &schema)
.with_child_sort(vec![vec!["a1", "b1", "c1"]], &schema2)
.with_child_sort(vec![vec!["a2", "b2", "c2"]], &schema3)
.with_expected_sort(vec![vec!["a", "b"]])
.run()
}
#[test]
fn test_union_equivalence_properties_multi_children_4() {
let schema = create_test_schema().unwrap();
let schema2 = append_fields(&schema, "1");
let schema3 = append_fields(&schema, "2");
UnionEquivalenceTest::new(&schema)
.with_child_sort(vec![vec!["a", "b"]], &schema)
.with_child_sort(vec![vec!["a1", "b1"]], &schema2)
.with_child_sort(vec![vec!["b2", "c2"]], &schema3)
.with_expected_sort(vec![])
.run()
}
#[test]
fn test_union_equivalence_properties_multi_children_5() {
let schema = create_test_schema().unwrap();
let schema2 = append_fields(&schema, "1");
UnionEquivalenceTest::new(&schema)
.with_child_sort(vec![vec!["a", "b"], vec!["c"]], &schema)
.with_child_sort(vec![vec!["a1", "b1"], vec!["c1"]], &schema2)
.with_expected_sort(vec![vec!["a", "b"], vec!["c"]])
.run()
}
#[test]
fn test_union_equivalence_properties_constants_common_constants() {
let schema = create_test_schema().unwrap();
UnionEquivalenceTest::new(&schema)
.with_child_sort_and_const_exprs(
vec![vec!["a"]],
vec!["b", "c"],
&schema,
)
.with_child_sort_and_const_exprs(
vec![vec!["b"]],
vec!["a", "c"],
&schema,
)
.with_expected_sort_and_const_exprs(
vec![vec!["a"], vec!["b"]],
vec!["c"],
)
.run()
}
#[test]
fn test_union_equivalence_properties_constants_prefix() {
let schema = create_test_schema().unwrap();
UnionEquivalenceTest::new(&schema)
.with_child_sort_and_const_exprs(
vec![vec!["a"]],
vec![],
&schema,
)
.with_child_sort_and_const_exprs(
vec![vec!["a", "b"]],
vec![],
&schema,
)
.with_expected_sort_and_const_exprs(
vec![vec!["a"]],
vec![],
)
.run()
}
#[test]
fn test_union_equivalence_properties_constants_asc_desc_mismatch() {
let schema = create_test_schema().unwrap();
UnionEquivalenceTest::new(&schema)
.with_child_sort_and_const_exprs(
vec![vec!["a"]],
vec![],
&schema,
)
.with_child_sort_and_const_exprs(
vec![vec!["a DESC"]],
vec![],
&schema,
)
.with_expected_sort_and_const_exprs(
vec![],
vec![],
)
.run()
}
#[test]
fn test_union_equivalence_properties_constants_different_schemas() {
let schema = create_test_schema().unwrap();
let schema2 = append_fields(&schema, "1");
UnionEquivalenceTest::new(&schema)
.with_child_sort_and_const_exprs(
vec![vec!["a"]],
vec![],
&schema,
)
.with_child_sort_and_const_exprs(
vec![vec!["a1", "b1"]],
vec![],
&schema2,
)
.with_expected_sort_and_const_exprs(
vec![vec!["a"]],
vec![],
)
.run()
}
#[test]
fn test_union_equivalence_properties_constants_fill_gaps() {
let schema = create_test_schema().unwrap();
UnionEquivalenceTest::new(&schema)
.with_child_sort_and_const_exprs(
vec![vec!["a", "c"]],
vec!["b"],
&schema,
)
.with_child_sort_and_const_exprs(
vec![vec!["b", "c"]],
vec!["a"],
&schema,
)
.with_expected_sort_and_const_exprs(
vec![vec!["a", "b", "c"], vec!["b", "a", "c"]],
vec![],
)
.run()
}
#[test]
fn test_union_equivalence_properties_constants_no_fill_gaps() {
let schema = create_test_schema().unwrap();
UnionEquivalenceTest::new(&schema)
.with_child_sort_and_const_exprs(
vec![vec!["a", "c"]],
vec!["d"],
&schema,
)
.with_child_sort_and_const_exprs(
vec![vec!["b", "c"]],
vec!["a"],
&schema,
)
.with_expected_sort_and_const_exprs(
vec![vec!["a"]],
vec![],
)
.run()
}
#[test]
fn test_union_equivalence_properties_constants_fill_some_gaps() {
let schema = create_test_schema().unwrap();
UnionEquivalenceTest::new(&schema)
.with_child_sort_and_const_exprs(
vec![vec!["c"]],
vec!["a", "b"],
&schema,
)
.with_child_sort_and_const_exprs(
vec![vec!["a DESC", "b"]],
vec![],
&schema,
)
.with_expected_sort_and_const_exprs(
vec![vec!["a DESC", "b"]],
vec![],
)
.run()
}
#[test]
fn test_union_equivalence_properties_constants_fill_gaps_non_symmetric() {
let schema = create_test_schema().unwrap();
UnionEquivalenceTest::new(&schema)
.with_child_sort_and_const_exprs(
vec![vec!["a", "c"]],
vec!["b"],
&schema,
)
.with_child_sort_and_const_exprs(
vec![vec!["b DESC", "c"]],
vec!["a"],
&schema,
)
.with_expected_sort_and_const_exprs(
vec![vec!["a", "b DESC", "c"], vec!["b DESC", "a", "c"]],
vec![],
)
.run()
}
#[test]
fn test_union_equivalence_properties_constants_gap_fill_symmetric() {
let schema = create_test_schema().unwrap();
UnionEquivalenceTest::new(&schema)
.with_child_sort_and_const_exprs(
vec![vec!["a", "b", "d"]],
vec!["c"],
&schema,
)
.with_child_sort_and_const_exprs(
vec![vec!["a", "c", "d"]],
vec!["b"],
&schema,
)
.with_expected_sort_and_const_exprs(
vec![vec!["a", "c", "b", "d"], vec!["a", "b", "c", "d"]],
vec![],
)
.run()
}
#[test]
fn test_union_equivalence_properties_constants_gap_fill_and_common() {
let schema = create_test_schema().unwrap();
UnionEquivalenceTest::new(&schema)
.with_child_sort_and_const_exprs(
vec![vec!["a DESC", "d"]],
vec!["b", "c"],
&schema,
)
.with_child_sort_and_const_exprs(
vec![vec!["a DESC", "c", "d"]],
vec!["b"],
&schema,
)
.with_expected_sort_and_const_exprs(
vec![vec!["a DESC", "c", "d"]],
vec!["b"],
)
.run()
}
#[test]
fn test_union_equivalence_properties_constants_middle_desc() {
let schema = create_test_schema().unwrap();
UnionEquivalenceTest::new(&schema)
.with_child_sort_and_const_exprs(
vec![vec!["a", "b DESC", "d"]],
vec!["c"],
&schema,
)
.with_child_sort_and_const_exprs(
vec![vec!["a", "c", "d"]],
vec!["b"],
&schema,
)
.with_expected_sort_and_const_exprs(
vec![vec!["a", "c", "b DESC", "d"], vec!["a", "b DESC", "c", "d"]],
vec![],
)
.run()
}
#[derive(Debug)]
struct UnionEquivalenceTest {
output_schema: SchemaRef,
child_properties: Vec<EquivalenceProperties>,
expected_properties: Option<EquivalenceProperties>,
}
impl UnionEquivalenceTest {
fn new(output_schema: &SchemaRef) -> Self {
Self {
output_schema: Arc::clone(output_schema),
child_properties: vec![],
expected_properties: None,
}
}
fn with_child_sort(
mut self,
orderings: Vec<Vec<&str>>,
schema: &SchemaRef,
) -> Self {
let properties = self.make_props(orderings, vec![], schema);
self.child_properties.push(properties);
self
}
fn with_child_sort_and_const_exprs(
mut self,
orderings: Vec<Vec<&str>>,
constants: Vec<&str>,
schema: &SchemaRef,
) -> Self {
let properties = self.make_props(orderings, constants, schema);
self.child_properties.push(properties);
self
}
fn with_expected_sort(mut self, orderings: Vec<Vec<&str>>) -> Self {
let properties = self.make_props(orderings, vec![], &self.output_schema);
self.expected_properties = Some(properties);
self
}
fn with_expected_sort_and_const_exprs(
mut self,
orderings: Vec<Vec<&str>>,
constants: Vec<&str>,
) -> Self {
let properties = self.make_props(orderings, constants, &self.output_schema);
self.expected_properties = Some(properties);
self
}
fn run(self) {
let Self {
output_schema,
child_properties,
expected_properties,
} = self;
let expected_properties =
expected_properties.expect("expected_properties not set");
for child_properties in child_properties
.iter()
.cloned()
.permutations(child_properties.len())
{
println!("--- permutation ---");
for c in &child_properties {
println!("{c}");
}
let actual_properties =
calculate_union(child_properties, Arc::clone(&output_schema))
.expect("failed to calculate union equivalence properties");
assert_eq_properties_same(
&actual_properties,
&expected_properties,
format!(
"expected: {expected_properties:?}\nactual: {actual_properties:?}"
),
);
}
}
fn make_props(
&self,
orderings: Vec<Vec<&str>>,
constants: Vec<&str>,
schema: &SchemaRef,
) -> EquivalenceProperties {
let orderings = orderings
.iter()
.map(|ordering| {
ordering
.iter()
.map(|name| parse_sort_expr(name, schema))
.collect::<LexOrdering>()
})
.collect::<Vec<_>>();
let constants = constants
.iter()
.map(|col_name| ConstExpr::new(col(col_name, schema).unwrap()))
.collect::<Vec<_>>();
EquivalenceProperties::new_with_orderings(Arc::clone(schema), &orderings)
.with_constants(constants)
}
}
fn assert_eq_properties_same(
lhs: &EquivalenceProperties,
rhs: &EquivalenceProperties,
err_msg: String,
) {
let lhs_constants = lhs.constants();
let rhs_constants = rhs.constants();
for rhs_constant in rhs_constants {
assert!(
const_exprs_contains(lhs_constants, rhs_constant.expr()),
"{err_msg}\nlhs: {lhs}\nrhs: {rhs}"
);
}
assert_eq!(
lhs_constants.len(),
rhs_constants.len(),
"{err_msg}\nlhs: {lhs}\nrhs: {rhs}"
);
let lhs_orderings = lhs.oeq_class();
let rhs_orderings = &rhs.oeq_class.orderings;
for rhs_ordering in rhs_orderings {
assert!(
lhs_orderings.contains(rhs_ordering),
"{err_msg}\nlhs: {lhs}\nrhs: {rhs}"
);
}
assert_eq!(
lhs_orderings.len(),
rhs_orderings.len(),
"{err_msg}\nlhs: {lhs}\nrhs: {rhs}"
);
}
fn parse_sort_expr(name: &str, schema: &SchemaRef) -> PhysicalSortExpr {
let mut parts = name.split_whitespace();
let name = parts.next().expect("empty sort expression");
let mut sort_expr = PhysicalSortExpr::new(
col(name, schema).expect("invalid column name"),
SortOptions::default(),
);
if let Some(options) = parts.next() {
sort_expr = match options {
"ASC" => sort_expr.asc(),
"DESC" => sort_expr.desc(),
_ => panic!(
"unknown sort options. Expected 'ASC' or 'DESC', got {}",
options
),
}
}
assert!(
parts.next().is_none(),
"unexpected tokens in column name. Expected 'name' / 'name ASC' / 'name DESC' but got '{name}'"
);
sort_expr
}
}