blaginin commented on code in PR #27439:
URL: https://github.com/apache/datafusion/pull/27439#discussion_r4575250320
##########
datafusion/physical-plan/src/windows/mod.rs:
##########
@@ -374,206 +374,164 @@ pub(crate) fn window_equivalence_properties(
input: &Arc<dyn ExecutionPlan>,
window_exprs: &[Arc<dyn WindowExpr>],
) -> Result<EquivalenceProperties> {
- // We need to update the schema, so we can't directly use input's
equivalence
- // properties.
let mut window_eq_properties =
EquivalenceProperties::new(Arc::clone(schema))
.extend(input.equivalence_properties().clone())?;
-
- let window_schema_len = schema.fields.len();
- let input_schema_len = window_schema_len - window_exprs.len();
- let window_expr_indices =
(input_schema_len..window_schema_len).collect::<Vec<_>>();
+ let input_schema_len = input.schema().fields().len();
for (i, expr) in window_exprs.iter().enumerate() {
let partitioning_exprs = expr.partition_by();
- let no_partitioning = partitioning_exprs.is_empty();
-
- // Find "one" valid ordering for partition columns to avoid
exponential complexity.
- // see https://github.com/apache/datafusion/issues/17401
- let mut all_satisfied_lexs = vec![];
- let mut candidate_ordering =
Vec::with_capacity(partitioning_exprs.len());
-
- for partition_expr in partitioning_exprs.iter() {
- let sort_options =
- sort_options_resolving_constant(Arc::clone(partition_expr),
true);
-
- // Try each sort option and pick the first one that works
- let mut found = false;
- for sort_expr in sort_options.into_iter() {
- candidate_ordering.push(sort_expr);
- if let Some(lex) = LexOrdering::new(candidate_ordering.clone())
- && window_eq_properties.ordering_satisfy(lex)?
- {
- found = true;
- break;
- }
- // This option didn't work, remove it and try the next one
- candidate_ordering.pop();
- }
- // If no sort option works for this column, we can't build a valid
ordering
- if !found {
- candidate_ordering.clear();
- break;
- }
+ let (partition_ordering, _) =
+ window_eq_properties.find_longest_permutation(partitioning_exprs)?;
+ if partition_ordering.len() != partitioning_exprs.len() {
+ continue;
}
- // If we successfully built an ordering for all columns, use it
- // When there are no partition expressions, candidate_ordering will be
empty and won't be added
- if candidate_ordering.len() == partitioning_exprs.len()
- && let Some(lex) = LexOrdering::new(candidate_ordering)
+ let window_col_idx = i + input_schema_len;
+ let frame = expr.get_window_frame();
+ let std_expr = expr.as_any().downcast_ref::<StandardWindowExpr>();
+ let aggregate = aggregate_and_filter(expr.as_ref());
+ if frame.is_ever_expanding()
+ && frame.end_bound.is_unbounded()
+ && (aggregate.is_some()
+ || std_expr.is_some_and(|expr| {
+ expr.get_standard_func_expr()
+ .as_any()
+ .downcast_ref::<WindowUDFExpr>()
+ .is_some_and(|expr|
expr.fun().inner().is_constant_for_frame())
+ }))
{
- all_satisfied_lexs.push(lex);
- }
- // If there is a partitioning, and no possible ordering cannot satisfy
- // the input plan's orderings, then we cannot further introduce any
- // new orderings for the window plan.
- if !no_partitioning && all_satisfied_lexs.is_empty() {
- return Ok(window_eq_properties);
- } else if let Some(std_expr) =
expr.as_any().downcast_ref::<StandardWindowExpr>()
+ add_constant_window_orderings(
+ &mut window_eq_properties,
+ Arc::new(Column::new(expr.name(), window_col_idx)),
+ &partition_ordering,
+ )?;
+ } else if let Some((aggregate, has_filter)) = aggregate
+ && let Some(growing) =
aggregate.get_result_ordering(window_col_idx)
{
+ add_aggregate_window_orderings(
+ &mut window_eq_properties,
+ aggregate,
+ frame,
+ has_filter,
+ growing,
+ &partition_ordering,
+ )?;
+ } else if let Some(std_expr) = std_expr {
std_expr.add_equal_orderings(&mut window_eq_properties)?;
- } else if let Some(plain_expr) =
- expr.as_any().downcast_ref::<PlainAggregateWindowExpr>()
- {
- // We are dealing with plain window frames; i.e. frames having an
- // unbounded starting point.
- // First, check if the frame covers the whole table:
- if plain_expr.get_window_frame().end_bound.is_unbounded() {
- let window_col =
- Arc::new(Column::new(expr.name(), i + input_schema_len))
as _;
- if no_partitioning {
- // Window function has a constant result across the table:
- window_eq_properties
-
.add_constants(std::iter::once(ConstExpr::from(window_col)))?
- } else {
- // Window function results in a partial constant value in
- // some ordering. Adjust the ordering equivalences
accordingly:
- let new_lexs =
all_satisfied_lexs.into_iter().flat_map(|lex| {
- let new_partial_consts =
sort_options_resolving_constant(
- Arc::clone(&window_col),
- false,
- );
-
- new_partial_consts.into_iter().map(move |partial| {
- let mut existing = lex.clone();
- existing.push(partial);
- existing
- })
- });
- window_eq_properties.add_orderings(new_lexs);
- }
- } else {
- // The window frame is ever expanding, so set monotonicity
comes
- // into play.
- plain_expr.add_equal_orderings(
- &mut window_eq_properties,
- window_expr_indices[i],
- )?;
- }
- } else if let Some(sliding_expr) =
- expr.as_any().downcast_ref::<SlidingAggregateWindowExpr>()
- {
- // We are dealing with sliding window frames; i.e. frames having an
- // advancing starting point. If we have a set-monotonic expression,
- // we might be able to leverage this property.
- let set_monotonicity =
sliding_expr.get_aggregate_expr().set_monotonicity();
- if set_monotonicity.ne(&SetMonotonicity::NotMonotonic) {
- // If the window frame is ever-receding, and we have set
- // monotonicity, we can utilize it to introduce new orderings.
- let frame = sliding_expr.get_window_frame();
- if frame.end_bound.is_unbounded() {
- let increasing =
set_monotonicity.eq(&SetMonotonicity::Increasing);
- let window_col = Column::new(expr.name(), i +
input_schema_len);
- if no_partitioning {
- // Reverse set-monotonic cases with no partitioning:
-
window_eq_properties.add_ordering([PhysicalSortExpr::new(
- Arc::new(window_col),
- SortOptions::new(increasing, true),
- )]);
- } else {
- // Reverse set-monotonic cases for all orderings:
- for mut lex in all_satisfied_lexs.into_iter() {
- lex.push(PhysicalSortExpr::new(
- Arc::new(window_col.clone()),
- SortOptions::new(increasing, true),
- ));
- window_eq_properties.add_ordering(lex);
- }
- }
- }
- // If we ensure that the elements entering the frame is greater
- // than the ones leaving, and we have increasing
set-monotonicity,
- // then the window function result will be increasing. However,
- // we also need to check if the frame is causal. If not, we
cannot
- // utilize set-monotonicity since the set shrinks as the frame
- // boundary starts "touching" the end of the table.
- else if frame.is_causal() {
- // Find one valid ordering for aggregate arguments instead
of
- // checking all combinations
- let aggregate_exprs =
sliding_expr.get_aggregate_expr().expressions();
- let mut candidate_order =
Vec::with_capacity(aggregate_exprs.len());
- let mut asc = false;
-
- for (idx, expr) in aggregate_exprs.iter().enumerate() {
- let mut found = false;
- let sort_options =
- sort_options_resolving_constant(Arc::clone(expr),
false);
-
- // Try each option and pick the first that works
- for sort_expr in sort_options.into_iter() {
- let is_asc = !sort_expr.options.descending;
- candidate_order.push(sort_expr);
-
- if let Some(lex) =
LexOrdering::new(candidate_order.clone())
- && window_eq_properties.ordering_satisfy(lex)?
- {
- if idx == 0 {
- // The first column's ordering direction
determines the overall
- // monotonicity behavior of the window
result.
- // - If the aggregate has increasing set
monotonicity (e.g., MAX, COUNT)
- // and the first arg is ascending, the
window result is increasing
- // - If the aggregate has decreasing set
monotonicity (e.g., MIN)
- // and the first arg is ascending, the
window result is also increasing
- // This flag is used to determine the
final window column ordering.
- asc = is_asc;
- }
- found = true;
- break;
- }
- // This option didn't work, remove it and try the
next one
- candidate_order.pop();
- }
-
- // If we couldn't extend the ordering, stop trying
- if !found {
- break;
- }
- }
+ }
+ }
+ Ok(window_eq_properties)
+}
- // Check if we successfully built a complete ordering
- let satisfied = candidate_order.len() ==
aggregate_exprs.len()
- && !aggregate_exprs.is_empty();
-
- if satisfied {
- let increasing =
- set_monotonicity.eq(&SetMonotonicity::Increasing);
- let window_col = Column::new(expr.name(), i +
input_schema_len);
- if increasing && (asc || no_partitioning) {
-
window_eq_properties.add_ordering([PhysicalSortExpr::new(
- Arc::new(window_col),
- SortOptions::new(false, false),
- )]);
- } else if !increasing && (!asc || no_partitioning) {
-
window_eq_properties.add_ordering([PhysicalSortExpr::new(
- Arc::new(window_col),
- SortOptions::new(true, false),
- )]);
- }
- }
- }
+fn aggregate_and_filter(expr: &dyn WindowExpr) ->
Option<(&AggregateFunctionExpr, bool)> {
+ if let Some(expr) =
expr.as_any().downcast_ref::<PlainAggregateWindowExpr>() {
+ Some((expr.get_aggregate_expr(), expr.has_filter()))
+ } else {
+ expr.as_any()
+ .downcast_ref::<SlidingAggregateWindowExpr>()
+ .map(|expr| (expr.get_aggregate_expr(), expr.has_filter()))
+ }
+}
+
+fn add_constant_window_orderings(
+ eq_properties: &mut EquivalenceProperties,
+ window_col: Arc<dyn PhysicalExpr>,
+ partition_ordering: &[PhysicalSortExpr],
+) -> Result<()> {
+ if partition_ordering.is_empty() {
+ eq_properties.add_constants([ConstExpr::from(window_col)])?;
+ } else {
+ for sort_expr in sort_options_resolving_constant(window_col) {
+ let mut ordering = partition_ordering.to_vec();
+ ordering.push(sort_expr);
+ eq_properties.add_ordering(ordering);
+ }
+ }
+ Ok(())
+}
+
+fn add_window_ordering(
+ eq_properties: &mut EquivalenceProperties,
+ mut sort_expr: PhysicalSortExpr,
+ partition_ordering: &[PhysicalSortExpr],
+ result_never_null: bool,
+) {
+ let mut ordering = partition_ordering.to_vec();
+ ordering.push(sort_expr.clone());
+ eq_properties.add_ordering(ordering);
+ if result_never_null {
+ sort_expr.options.nulls_first = !sort_expr.options.nulls_first;
+ let mut ordering = partition_ordering.to_vec();
+ ordering.push(sort_expr);
+ eq_properties.add_ordering(ordering);
+ }
+}
+
+fn frame_contains_current_row(frame: &WindowFrame) -> bool {
+ matches!(
+ frame.start_bound,
+ WindowFrameBound::Preceding(_) | WindowFrameBound::CurrentRow
+ ) && matches!(
+ frame.end_bound,
+ WindowFrameBound::CurrentRow | WindowFrameBound::Following(_)
+ )
+}
+
+fn add_aggregate_window_orderings(
+ eq_properties: &mut EquivalenceProperties,
+ aggregate: &AggregateFunctionExpr,
+ frame: &WindowFrame,
+ has_filter: bool,
+ growing: PhysicalSortExpr,
+ partition_ordering: &[PhysicalSortExpr],
+) -> Result<()> {
+ let arguments = aggregate.expressions();
+ let frame_never_empty = !has_filter && frame_contains_current_row(frame);
+ let mut result_never_null = !aggregate.is_nullable();
+ if !result_never_null && frame_never_empty {
+ result_never_null = !arguments.is_empty();
+ for argument in &arguments {
+ result_never_null &= !argument.nullable(eq_properties.schema())?;
+ }
+ }
+
+ match (frame.is_ever_expanding(), frame.end_bound.is_unbounded()) {
+ (true, false) => add_window_ordering(
+ eq_properties,
+ growing.clone(),
+ partition_ordering,
+ result_never_null,
+ ),
+ (false, true) => add_window_ordering(
+ eq_properties,
+ growing.reverse(),
+ partition_ordering,
+ result_never_null,
+ ),
+ _ => {}
Review Comment:
@itang06 i think this partially covers #27267 if you want to review this
bit. also happy to remove it from here so you can merge your change 🙂
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
