avantgardnerio commented on code in PR #2885:
URL: https://github.com/apache/arrow-datafusion/pull/2885#discussion_r926103558
##########
datafusion/optimizer/src/utils.rs:
##########
@@ -82,6 +119,199 @@ pub fn add_filter(plan: LogicalPlan, predicates: &[&Expr])
-> LogicalPlan {
})
}
+/// Looks for correlating expressions: equality expressions with one field
from the subquery, and
+/// one not in the subquery (closed upon from outer scope)
+///
+/// # Arguments
+///
+/// * `exprs` - List of expressions that may or may not be joins
+/// * `fields` - HashSet of fully qualified (table.col) fields in subquery
schema
+///
+/// # Return value
+///
+/// Tuple of (expressions containing joins, remaining non-join expressions)
+pub fn find_join_exprs(
+ exprs: Vec<&Expr>,
+ schema: &DFSchemaRef,
+) -> Result<(Vec<Expr>, Vec<Expr>)> {
+ let fields: HashSet<_> = schema
+ .fields()
+ .iter()
+ .map(|it| it.qualified_name())
+ .collect();
+
+ let mut joins = vec![];
+ let mut others = vec![];
+ for filter in exprs.iter() {
+ let (left, op, right) = match filter {
+ Expr::BinaryExpr { left, op, right } => (*left.clone(), *op,
*right.clone()),
+ _ => {
+ others.push((*filter).clone());
+ continue;
+ }
+ };
+ let left = match left {
+ Expr::Column(c) => c,
+ _ => {
+ others.push((*filter).clone());
+ continue;
+ }
+ };
+ let right = match right {
+ Expr::Column(c) => c,
+ _ => {
+ others.push((*filter).clone());
+ continue;
+ }
+ };
+ if fields.contains(&left.flat_name()) &&
fields.contains(&right.flat_name()) {
+ others.push((*filter).clone());
+ continue; // both columns present (none closed-upon)
+ }
+ if !fields.contains(&left.flat_name()) &&
!fields.contains(&right.flat_name()) {
+ others.push((*filter).clone());
+ continue; // neither column present (syntax error?)
+ }
+ match op {
+ Operator::Eq => {}
+ Operator::NotEq => {}
+ _ => {
+ plan_err!(format!("can't optimize {} column comparison", op))?;
+ }
+ }
+
+ joins.push((*filter).clone())
+ }
+
+ Ok((joins, others))
+}
+
+/// Extracts correlating columns from expressions
+///
+/// # Arguments
+///
+/// * `exprs` - List of expressions that correlate a subquery to an outer scope
+/// * `fields` - HashSet of fully qualified (table.col) fields in subquery
schema
+///
+/// # Return value
+///
+/// Tuple of tuples ((outer-scope cols, subquery cols), non-equal expressions)
+pub fn exprs_to_join_cols(
+ exprs: &[Expr],
+ schema: &DFSchemaRef,
+ include_negated: bool,
+) -> Result<(Vec<Column>, Vec<Column>, Option<Expr>)> {
+ let fields: HashSet<_> = schema
+ .fields()
+ .iter()
+ .map(|it| it.qualified_name())
+ .collect();
+
+ let mut joins: Vec<(String, String)> = vec![];
+ let mut others: Vec<Expr> = vec![];
+ for filter in exprs.iter() {
+ let (left, op, right) = match filter {
+ Expr::BinaryExpr { left, op, right } => (*left.clone(), *op,
*right.clone()),
+ _ => plan_err!("Invalid expression!".to_string())?,
+ };
+ match op {
+ Operator::Eq => {}
+ Operator::NotEq => {
+ if !include_negated {
+ others.push((*filter).clone());
+ continue;
+ }
+ }
+ _ => plan_err!("Invalid expression!".to_string())?,
+ }
+ let left = match left {
+ Expr::Column(c) => c,
+ _ => plan_err!("Invalid expression!".to_string())?,
+ };
+ let right = match right {
+ Expr::Column(c) => c,
+ _ => plan_err!("Invalid expression!".to_string())?,
+ };
+ let sorted = if fields.contains(&left.flat_name()) {
+ (right.flat_name(), left.flat_name())
+ } else {
+ (left.flat_name(), right.flat_name())
+ };
+ joins.push(sorted);
+ }
+
+ let (left_cols, right_cols): (Vec<_>, Vec<_>) = joins
+ .into_iter()
+ .map(|(l, r)| (Column::from(l.as_str()), Column::from(r.as_str())))
+ .unzip();
+ let pred = combine_filters(&others);
+
+ Ok((left_cols, right_cols, pred))
+}
+
+/// Returns the first (and only) element in a slice, or an error
+///
+/// # Arguments
+///
+/// * `slice` - The slice to extract from
+///
+/// # Return value
+///
+/// The first element, or an error
+pub fn only_or_err<T>(slice: &[T]) -> Result<&T> {
+ match slice {
+ [it] => Ok(it),
+ [] => Err(DataFusionError::Plan("Empty slice!".to_string())),
+ _ => Err(DataFusionError::Plan(
+ "More than one item in slice!".to_string(),
+ )),
+ }
+}
+
+/// Merge and deduplicate two Column slices
+///
+/// # Arguments
+///
+/// * `a` - A slice of Columns
+/// * `b` - A slice of Columns
+///
+/// # Return value
+///
+/// The deduplicated union of the two slices
+pub fn merge_cols(a: &[Column], b: &[Column]) -> Vec<Column> {
+ let a: Vec<_> = a.iter().map(|it| it.flat_name()).collect();
Review Comment:
Updated. It's efficient now, but ugly.
--
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]
