Skip to content

Commit

Permalink
Fix comment
Browse files Browse the repository at this point in the history
  • Loading branch information
xinlifoobar committed Jul 23, 2024
1 parent e8b097b commit 86cfd89
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 13 deletions.
20 changes: 20 additions & 0 deletions datafusion/expr/src/expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1340,6 +1340,9 @@ impl Expr {
///
/// returns `None` if the expression is not a `Column`
///
/// Note: None may be returned for expressions that are not `Column` but
/// are convertible to `Column` such as `Cast` expressions.
///
/// Example
/// ```
/// # use datafusion_common::Column;
Expand All @@ -1358,6 +1361,23 @@ impl Expr {
}
}

/// Returns the inner `Column` if any. This is a specialized version of
/// [`Self::try_as_col`] that take Cast expressions into account when the
/// expression is as on condition for joins.
///
/// Called this method when you are sure that the expression is a `Column`
/// or a `Cast` expression that wraps a `Column`.
pub fn get_as_join_column(&self) -> Option<&Column> {
match self {
Expr::Column(c) => Some(c),
Expr::Cast(Cast { expr, .. }) => match &**expr {
Expr::Column(c) => Some(c),
_ => None,
},
_ => None,
}
}

/// Return all referenced columns of this expression.
#[deprecated(since = "40.0.0", note = "use Expr::column_refs instead")]
pub fn to_columns(&self) -> Result<HashSet<Column>> {
Expand Down
26 changes: 13 additions & 13 deletions datafusion/expr/src/logical_plan/plan.rs
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ use crate::utils::{
grouping_set_expr_count, grouping_set_to_exprlist, split_conjunction,
};
use crate::{
build_join_schema, expr_vec_fmt, BinaryExpr, BuiltInWindowFunction, Cast,
build_join_schema, expr_vec_fmt, BinaryExpr, BuiltInWindowFunction,
CreateMemoryTable, CreateView, Expr, ExprSchemable, LogicalPlanBuilder, Operator,
TableProviderFilterPushDown, TableSource, WindowFunctionDefinition,
};
Expand Down Expand Up @@ -496,10 +496,18 @@ impl LogicalPlan {
// The join keys in using-join must be columns.
let columns =
on.iter().try_fold(HashSet::new(), |mut accumu, (l, r)| {
let l = Self::as_col(l.clone())?;
let r = Self::as_col(r.clone())?;
accumu.insert(l);
accumu.insert(r);
let Some(l) = l.get_as_join_column() else {
return internal_err!(
"Invalid join key. Expected column, found {l:?}"
);
};
let Some(r) = r.get_as_join_column() else {
return internal_err!(
"Invalid join key. Expected column, found {r:?}"
);
};
accumu.insert(l.to_owned());
accumu.insert(r.to_owned());
Result::<_, DataFusionError>::Ok(accumu)
})?;
using_columns.push(columns);
Expand All @@ -510,14 +518,6 @@ impl LogicalPlan {
Ok(using_columns)
}

fn as_col(expr: Expr) -> Result<Column> {
match expr {
Expr::Column(c) => Ok(c),
Expr::Cast(Cast { expr, .. }) => Self::as_col(*expr),
_ => internal_err!("Invalid join key. Expected column, found {expr:?}"),
}
}

/// returns the first output expression of this `LogicalPlan` node.
pub fn head_output_expr(&self) -> Result<Option<Expr>> {
match self {
Expand Down

0 comments on commit 86cfd89

Please sign in to comment.