Skip to content

Commit 380d843

Browse files
authored
Improve documentation and add Display impl to EquivalenceProperties (#12590)
1 parent 6546479 commit 380d843

File tree

8 files changed

+200
-33
lines changed

8 files changed

+200
-33
lines changed

datafusion/core/tests/fuzz_cases/sort_preserving_repartition_fuzz.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,7 @@ mod sp_repartition_fuzz_tests {
8080
// Define a and f are aliases
8181
eq_properties.add_equal_conditions(col_a, col_f)?;
8282
// Column e has constant value.
83-
eq_properties = eq_properties.add_constants([ConstExpr::from(col_e)]);
83+
eq_properties = eq_properties.with_constants([ConstExpr::from(col_e)]);
8484

8585
// Randomly order columns for sorting
8686
let mut rng = StdRng::seed_from_u64(seed);

datafusion/physical-expr-common/src/physical_expr.rs

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
// under the License.
1717

1818
use std::any::Any;
19-
use std::fmt::{Debug, Display};
19+
use std::fmt::{Debug, Display, Formatter};
2020
use std::hash::{Hash, Hasher};
2121
use std::sync::Arc;
2222

@@ -223,3 +223,25 @@ pub fn down_cast_any_ref(any: &dyn Any) -> &dyn Any {
223223
any
224224
}
225225
}
226+
227+
/// Returns [`Display`] able a list of [`PhysicalExpr`]
228+
///
229+
/// Example output: `[a + 1, b]`
230+
pub fn format_physical_expr_list(exprs: &[Arc<dyn PhysicalExpr>]) -> impl Display + '_ {
231+
struct DisplayWrapper<'a>(&'a [Arc<dyn PhysicalExpr>]);
232+
impl<'a> Display for DisplayWrapper<'a> {
233+
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
234+
let mut iter = self.0.iter();
235+
write!(f, "[")?;
236+
if let Some(expr) = iter.next() {
237+
write!(f, "{}", expr)?;
238+
}
239+
for expr in iter {
240+
write!(f, ", {}", expr)?;
241+
}
242+
write!(f, "]")?;
243+
Ok(())
244+
}
245+
}
246+
DisplayWrapper(exprs)
247+
}

datafusion/physical-expr/src/equivalence/class.rs

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18+
use std::fmt::Display;
1819
use std::sync::Arc;
1920

2021
use super::{add_offset_to_expr, collapse_lex_req, ProjectionMapping};
@@ -27,6 +28,7 @@ use crate::{
2728

2829
use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
2930
use datafusion_common::JoinType;
31+
use datafusion_physical_expr_common::physical_expr::format_physical_expr_list;
3032

3133
#[derive(Debug, Clone)]
3234
/// A structure representing a expression known to be constant in a physical execution plan.
@@ -101,6 +103,19 @@ impl ConstExpr {
101103
}
102104
}
103105

106+
/// Display implementation for `ConstExpr`
107+
///
108+
/// Example `c` or `c(across_partitions)`
109+
impl Display for ConstExpr {
110+
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
111+
write!(f, "{}", self.expr)?;
112+
if self.across_partitions {
113+
write!(f, "(across_partitions)")?;
114+
}
115+
Ok(())
116+
}
117+
}
118+
104119
impl From<Arc<dyn PhysicalExpr>> for ConstExpr {
105120
fn from(expr: Arc<dyn PhysicalExpr>) -> Self {
106121
Self::new(expr)
@@ -224,6 +239,12 @@ impl EquivalenceClass {
224239
}
225240
}
226241

242+
impl Display for EquivalenceClass {
243+
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
244+
write!(f, "[{}]", format_physical_expr_list(&self.exprs))
245+
}
246+
}
247+
227248
/// An `EquivalenceGroup` is a collection of `EquivalenceClass`es where each
228249
/// class represents a distinct equivalence class in a relation.
229250
#[derive(Debug, Clone)]
@@ -575,6 +596,20 @@ impl EquivalenceGroup {
575596
}
576597
}
577598

599+
impl Display for EquivalenceGroup {
600+
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
601+
write!(f, "[")?;
602+
let mut iter = self.iter();
603+
if let Some(cls) = iter.next() {
604+
write!(f, "{}", cls)?;
605+
}
606+
for cls in iter {
607+
write!(f, ", {}", cls)?;
608+
}
609+
write!(f, "]")
610+
}
611+
}
612+
578613
#[cfg(test)]
579614
mod tests {
580615

datafusion/physical-expr/src/equivalence/mod.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -207,7 +207,7 @@ mod tests {
207207
// Define a and f are aliases
208208
eq_properties.add_equal_conditions(col_a, col_f)?;
209209
// Column e has constant value.
210-
eq_properties = eq_properties.add_constants([ConstExpr::from(col_e)]);
210+
eq_properties = eq_properties.with_constants([ConstExpr::from(col_e)]);
211211

212212
// Randomly order columns for sorting
213213
let mut rng = StdRng::seed_from_u64(seed);

datafusion/physical-expr/src/equivalence/ordering.rs

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,13 +15,13 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18+
use std::fmt::Display;
1819
use std::hash::Hash;
1920
use std::sync::Arc;
2021

21-
use arrow_schema::SortOptions;
22-
2322
use crate::equivalence::add_offset_to_expr;
2423
use crate::{LexOrdering, PhysicalExpr, PhysicalSortExpr};
24+
use arrow_schema::SortOptions;
2525

2626
/// An `OrderingEquivalenceClass` object keeps track of different alternative
2727
/// orderings than can describe a schema. For example, consider the following table:
@@ -104,6 +104,11 @@ impl OrderingEquivalenceClass {
104104
self.remove_redundant_entries();
105105
}
106106

107+
/// Adds a single ordering to the existing ordering equivalence class.
108+
pub fn add_new_ordering(&mut self, ordering: LexOrdering) {
109+
self.add_new_orderings([ordering]);
110+
}
111+
107112
/// Removes redundant orderings from this equivalence class. For instance,
108113
/// if we already have the ordering `[a ASC, b ASC, c DESC]`, then there is
109114
/// no need to keep ordering `[a ASC, b ASC]` in the state.
@@ -219,6 +224,21 @@ fn resolve_overlap(orderings: &mut [LexOrdering], idx: usize, pre_idx: usize) ->
219224
false
220225
}
221226

227+
impl Display for OrderingEquivalenceClass {
228+
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
229+
write!(f, "[")?;
230+
let mut iter = self.orderings.iter();
231+
if let Some(ordering) = iter.next() {
232+
write!(f, "{}", PhysicalSortExpr::format_list(ordering))?;
233+
}
234+
for ordering in iter {
235+
write!(f, "{}", PhysicalSortExpr::format_list(ordering))?;
236+
}
237+
write!(f, "]")?;
238+
Ok(())
239+
}
240+
}
241+
222242
#[cfg(test)]
223243
mod tests {
224244
use std::sync::Arc;
@@ -559,7 +579,7 @@ mod tests {
559579
let constants = constants
560580
.into_iter()
561581
.map(|expr| ConstExpr::from(expr).with_across_partitions(true));
562-
eq_properties = eq_properties.add_constants(constants);
582+
eq_properties = eq_properties.with_constants(constants);
563583

564584
let reqs = convert_to_sort_exprs(&reqs);
565585
assert_eq!(

0 commit comments

Comments
 (0)