gene-bordegaray commented on code in PR #23184: URL: https://github.com/apache/datafusion/pull/23184#discussion_r3539187482
########## datafusion/physical-plan/src/distribution_requirements.rs: ########## @@ -0,0 +1,480 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Input distribution requirements for physical execution plans. + +use std::sync::Arc; + +use datafusion_common::{Result, internal_err}; +use datafusion_physical_expr::{ + Distribution, EquivalenceProperties, Partitioning, PartitioningSatisfaction, + PhysicalExpr, physical_exprs_equal, +}; + +use crate::execution_plan::{ExecutionPlan, ExecutionPlanProperties, InvariantLevel}; + +/// Distribution requirements for an [`ExecutionPlan`]'s inputs. +/// +/// [`RequiredInputDistributions`] describes what distribution an operator +/// requires from each child. +/// +/// - [`Self::new`] describes independent per-child requirements. +/// - [`Self::co_partitioned`] additionally requires child partitions with the +/// same index to cover compatible key ranges. +/// +/// For a single-input aggregate: +/// +/// ```text +/// AggregateExec +/// child 0 requirement: KeyPartitioned(group_exprs) +/// ``` +/// +/// each input partition can aggregate its own key domain independently. +/// +/// For a partitioned join: +/// +/// ```text +/// HashJoinExec +/// child 0 requirement: KeyPartitioned(left_keys) +/// child 1 requirement: KeyPartitioned(right_keys) +/// +/// partition 0: join(left partition 0, right partition 0) +/// partition 1: join(left partition 1, right partition 1) +/// partition 2: join(left partition 2, right partition 2) +/// ``` +/// +/// each child must satisfy its own key requirement. In addition, matching +/// partition indexes must be safe to process together. +#[non_exhaustive] +#[derive(Debug, Clone)] +pub struct RequiredInputDistributions { + /// Per-child distribution requirements, indexed by child position. + children: Vec<ChildDistributionRequirement>, + /// Child indexes that must also have compatible partition layouts. + co_partitioned: Option<Vec<usize>>, +} + +impl RequiredInputDistributions { + /// Create independent per-child requirements. + pub fn new(per_child: Vec<Distribution>) -> Self { + let children = per_child + .into_iter() + .map(|distribution| ChildDistributionRequirement { + distribution, + satisfaction: InputDistributionSatisfaction::Default, + }) + .collect(); + + Self { + children, + co_partitioned: None, + } + } + + /// Create a requirement that all children are co-partitioned. + /// + /// Each child must satisfy its own [`Distribution`]. Matching partition + /// indexes are processed together: + /// + /// ```text + /// left: Range(left.a ASC, split_points=[10, 20]) + /// right: Range(right.x ASC, split_points=[10, 20]) + /// + /// partition 0 from both sides contains keys before 10 + /// partition 1 from both sides contains keys in [10, 20) + /// partition 2 from both sides contains keys at/after 20 + /// ``` + /// + /// If the split points differ, partition `i` from one side no longer covers + /// the same key range as partition `i` from the other side. + pub fn co_partitioned(per_child: Vec<Distribution>) -> Self { + debug_assert!( + per_child.len() >= 2, + "co-partitioned distribution requirements need at least two children" + ); + let co_partitioned = (0..per_child.len()).collect(); + let mut result = Self::new(per_child); + result.co_partitioned = Some(co_partitioned); + result + } + + /// Return the per-child distribution requirements. + pub fn per_child_distributions( + &self, + ) -> impl ExactSizeIterator<Item = &Distribution> + '_ { + self.children.iter().map(|child| &child.distribution) + } + + /// Return the distribution requirement for a child. + pub fn child_distribution(&self, child_idx: usize) -> Option<&Distribution> { + self.children + .get(child_idx) + .map(|child| &child.distribution) + } + + /// Return the per-child distribution requirements. + /// + /// WARNING: This intentionally drops any grouped relationship. + pub fn into_per_child(self) -> Vec<Distribution> { + self.children + .into_iter() + .map(|child| child.distribution) + .collect() + } + + /// Returns how a child satisfies its distribution requirement using this + /// requirement set's satisfaction policy. + pub fn child_satisfaction( + &self, + child_idx: usize, + child: &dyn ExecutionPlan, + allow_subset: bool, + ) -> Result<PartitioningSatisfaction> { + let Some(requirement) = self.children.get(child_idx) else { + return internal_err!( + "missing distribution requirement for child {child_idx}" + ); + }; + + Ok(requirement.satisfaction.satisfaction( + child.output_partitioning(), + &requirement.distribution, + child.equivalence_properties(), + allow_subset, + )) + } + + /// Return child indexes whose co-partitioning requirements are + /// unsatisfied by the provided candidate children. + /// + /// Independent per-child requirements are intentionally ignored here, use + /// [`Self::child_satisfaction`] for those checks. An empty result means all + /// co-partitioning requirements are satisfied. + #[doc(hidden)] Review Comment: I am quite happy with the shape of this struct besides this method It is unfortunate that some optimizer specific thing needs to be public facing. This is done elsewhere like `Partitioning::satisfaction()` but this seems like users shouldn't really be touching. -- 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]
