Skip to content

Commit

Permalink
Update properties.rs
Browse files Browse the repository at this point in the history
  • Loading branch information
@berkaysynnada
berkaysynnada committed Dec 20, 2024
1 parent 4f8a206 commit 8d30c7b
Showing 1 changed file with 242 additions and 0 deletions.
242 changes: 242 additions & 0 deletions datafusion/physical-expr/src/equivalence/properties.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3891,4 +3891,246 @@ mod tests {

Ok(())
}

#[test]
fn test_with_reorder_constant_filtering() -> Result<()> {
let schema = create_test_schema()?;
let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema));

// Setup constant columns
let col_a = col("a", &schema)?;
let col_b = col("b", &schema)?;
eq_properties = eq_properties.with_constants([ConstExpr::from(&col_a)]);

let sort_exprs = LexOrdering::new(vec![
PhysicalSortExpr {
expr: Arc::clone(&col_a),
options: SortOptions::default(),
},
PhysicalSortExpr {
expr: Arc::clone(&col_b),
options: SortOptions::default(),
},
]);

let result = eq_properties.with_reorder(sort_exprs);

// Should only contain b since a is constant
assert_eq!(result.oeq_class().len(), 1);
assert_eq!(result.oeq_class().orderings[0].len(), 1);
assert!(result.oeq_class().orderings[0][0].expr.eq(&col_b));

Ok(())
}

#[test]
fn test_with_reorder_preserve_suffix() -> Result<()> {
let schema = create_test_schema()?;
let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema));

let col_a = col("a", &schema)?;
let col_b = col("b", &schema)?;
let col_c = col("c", &schema)?;

let asc = SortOptions::default();
let desc = SortOptions {
descending: true,
nulls_first: true,
};

// Initial ordering: [a ASC, b DESC, c ASC]
eq_properties.add_new_orderings([LexOrdering::new(vec![
PhysicalSortExpr {
expr: Arc::clone(&col_a),
options: asc,
},
PhysicalSortExpr {
expr: Arc::clone(&col_b),
options: desc,
},
PhysicalSortExpr {
expr: Arc::clone(&col_c),
options: asc,
},
])]);

// New ordering: [a ASC]
let new_order = LexOrdering::new(vec![PhysicalSortExpr {
expr: Arc::clone(&col_a),
options: asc,
}]);

let result = eq_properties.with_reorder(new_order);

// Should only contain [a ASC, b DESC, c ASC]
assert_eq!(result.oeq_class().len(), 1);
assert_eq!(result.oeq_class().orderings[0].len(), 3);
assert!(result.oeq_class().orderings[0][0].expr.eq(&col_a));
assert!(result.oeq_class().orderings[0][0].options.eq(&asc));
assert!(result.oeq_class().orderings[0][1].expr.eq(&col_b));
assert!(result.oeq_class().orderings[0][1].options.eq(&desc));
assert!(result.oeq_class().orderings[0][2].expr.eq(&col_c));
assert!(result.oeq_class().orderings[0][2].options.eq(&asc));

Ok(())
}

#[test]
fn test_with_reorder_equivalent_expressions() -> Result<()> {
let schema = create_test_schema()?;
let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema));

let col_a = col("a", &schema)?;
let col_b = col("b", &schema)?;
let col_c = col("c", &schema)?;

// Make a and b equivalent
eq_properties.add_equal_conditions(&col_a, &col_b)?;

let asc = SortOptions::default();

// Initial ordering: [a ASC, c ASC]
eq_properties.add_new_orderings([LexOrdering::new(vec![
PhysicalSortExpr {
expr: Arc::clone(&col_a),
options: asc,
},
PhysicalSortExpr {
expr: Arc::clone(&col_c),
options: asc,
},
])]);

// New ordering: [b ASC]
let new_order = LexOrdering::new(vec![PhysicalSortExpr {
expr: Arc::clone(&col_b),
options: asc,
}]);

let result = eq_properties.with_reorder(new_order);

// Should only contain [b ASC, c ASC]
assert_eq!(result.oeq_class().len(), 1);

// Verify orderings
assert_eq!(result.oeq_class().orderings[0].len(), 2);
assert!(result.oeq_class().orderings[0][0].expr.eq(&col_b));
assert!(result.oeq_class().orderings[0][0].options.eq(&asc));
assert!(result.oeq_class().orderings[0][1].expr.eq(&col_c));
assert!(result.oeq_class().orderings[0][1].options.eq(&asc));

Ok(())
}

#[test]
fn test_with_reorder_incompatible_prefix() -> Result<()> {
let schema = create_test_schema()?;
let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema));

let col_a = col("a", &schema)?;
let col_b = col("b", &schema)?;

let asc = SortOptions::default();
let desc = SortOptions {
descending: true,
nulls_first: true,
};

// Initial ordering: [a ASC, b DESC]
eq_properties.add_new_orderings([LexOrdering::new(vec![
PhysicalSortExpr {
expr: Arc::clone(&col_a),
options: asc,
},
PhysicalSortExpr {
expr: Arc::clone(&col_b),
options: desc,
},
])]);

// New ordering: [a DESC]
let new_order = LexOrdering::new(vec![PhysicalSortExpr {
expr: Arc::clone(&col_a),
options: desc,
}]);

let result = eq_properties.with_reorder(new_order.clone());

// Should only contain the new ordering since options don't match
assert_eq!(result.oeq_class().len(), 1);
assert_eq!(result.oeq_class().orderings[0], new_order);

Ok(())
}

#[test]
fn test_with_reorder_comprehensive() -> Result<()> {
let schema = create_test_schema()?;
let mut eq_properties = EquivalenceProperties::new(Arc::clone(&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 asc = SortOptions::default();

// Constants: c is constant
eq_properties = eq_properties.with_constants([ConstExpr::from(&col_c)]);

// Equality: b = d
eq_properties.add_equal_conditions(&col_b, &col_d)?;

// Orderings: [d ASC, a ASC], [e ASC]
eq_properties.add_new_orderings([
LexOrdering::new(vec![
PhysicalSortExpr {
expr: Arc::clone(&col_d),
options: asc,
},
PhysicalSortExpr {
expr: Arc::clone(&col_a),
options: asc,
},
]),
LexOrdering::new(vec![PhysicalSortExpr {
expr: Arc::clone(&col_e),
options: asc,
}]),
]);

// Initial ordering: [b ASC, c ASC]
let new_order = LexOrdering::new(vec![
PhysicalSortExpr {
expr: Arc::clone(&col_b),
options: asc,
},
PhysicalSortExpr {
expr: Arc::clone(&col_c),
options: asc,
},
]);

let result = eq_properties.with_reorder(new_order);

// Should preserve the original [d ASC, a ASC] ordering
assert_eq!(result.oeq_class().len(), 1);
let ordering = &result.oeq_class().orderings[0];
assert_eq!(ordering.len(), 2);

// First expression should be either b or d (they're equivalent)
assert!(
ordering[0].expr.eq(&col_b) || ordering[0].expr.eq(&col_d),
"Expected b or d as first expression, got {:?}",
ordering[0].expr
);
assert!(ordering[0].options.eq(&asc));

// Second expression should be a
assert!(ordering[1].expr.eq(&col_a));
assert!(ordering[1].options.eq(&asc));

Ok(())
}
}

0 comments on commit 8d30c7b

Please sign in to comment.