This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 53de994423 Add Expr::try_as_col, deprecate Expr::try_into_col (#10448)
53de994423 is described below
commit 53de994423fc85f655da232db7e807c2a38276ea
Author: Andrew Lamb <[email protected]>
AuthorDate: Mon May 13 09:10:13 2024 -0400
Add Expr::try_as_col, deprecate Expr::try_into_col (#10448)
---
datafusion/expr/src/expr.rs | 23 ++++++++++++++++++++++
datafusion/expr/src/expr_rewriter/mod.rs | 2 +-
datafusion/expr/src/logical_plan/builder.rs | 10 +++++++---
datafusion/expr/src/logical_plan/plan.rs | 14 +++++++++++--
datafusion/optimizer/src/push_down_filter.rs | 4 ++--
.../src/simplify_expressions/inlist_simplifier.rs | 2 +-
datafusion/proto/src/logical_plan/mod.rs | 12 ++++++++---
7 files changed, 55 insertions(+), 12 deletions(-)
diff --git a/datafusion/expr/src/expr.rs b/datafusion/expr/src/expr.rs
index f0f41a4c55..660a45c27a 100644
--- a/datafusion/expr/src/expr.rs
+++ b/datafusion/expr/src/expr.rs
@@ -1264,6 +1264,7 @@ impl Expr {
})
}
+ #[deprecated(since = "39.0.0", note = "use try_as_col instead")]
pub fn try_into_col(&self) -> Result<Column> {
match self {
Expr::Column(it) => Ok(it.clone()),
@@ -1271,6 +1272,28 @@ impl Expr {
}
}
+ /// Return a reference to the inner `Column` if any
+ ///
+ /// returns `None` if the expression is not a `Column`
+ ///
+ /// Example
+ /// ```
+ /// # use datafusion_common::Column;
+ /// use datafusion_expr::{col, Expr};
+ /// let expr = col("foo");
+ /// assert_eq!(expr.try_as_col(), Some(&Column::from("foo")));
+ ///
+ /// let expr = col("foo").alias("bar");
+ /// assert_eq!(expr.try_as_col(), None);
+ /// ```
+ pub fn try_as_col(&self) -> Option<&Column> {
+ if let Expr::Column(it) = self {
+ Some(it)
+ } else {
+ None
+ }
+ }
+
/// Return all referenced columns of this expression.
pub fn to_columns(&self) -> Result<HashSet<Column>> {
let mut using_columns = HashSet::new();
diff --git a/datafusion/expr/src/expr_rewriter/mod.rs
b/datafusion/expr/src/expr_rewriter/mod.rs
index 700dd560ec..1441374bdb 100644
--- a/datafusion/expr/src/expr_rewriter/mod.rs
+++ b/datafusion/expr/src/expr_rewriter/mod.rs
@@ -221,7 +221,7 @@ pub fn coerce_plan_expr_for_schema(
let exprs: Vec<Expr> =
plan.schema().iter().map(Expr::from).collect();
let new_exprs = coerce_exprs_for_schema(exprs, plan.schema(),
schema)?;
- let add_project = new_exprs.iter().any(|expr|
expr.try_into_col().is_err());
+ let add_project = new_exprs.iter().any(|expr|
expr.try_as_col().is_none());
if add_project {
let projection = Projection::try_new(new_exprs,
Arc::new(plan.clone()))?;
Ok(LogicalPlan::Projection(projection))
diff --git a/datafusion/expr/src/logical_plan/builder.rs
b/datafusion/expr/src/logical_plan/builder.rs
index 3f15b84784..2c6cfd8f9d 100644
--- a/datafusion/expr/src/logical_plan/builder.rs
+++ b/datafusion/expr/src/logical_plan/builder.rs
@@ -1489,7 +1489,7 @@ pub fn wrap_projection_for_join_if_necessary(
let mut projection = expand_wildcard(input_schema, &input, None)?;
let join_key_items = alias_join_keys
.iter()
- .flat_map(|expr| expr.try_into_col().is_err().then_some(expr))
+ .flat_map(|expr| expr.try_as_col().is_none().then_some(expr))
.cloned()
.collect::<HashSet<Expr>>();
projection.extend(join_key_items);
@@ -1504,8 +1504,12 @@ pub fn wrap_projection_for_join_if_necessary(
let join_on = alias_join_keys
.into_iter()
.map(|key| {
- key.try_into_col()
- .or_else(|_| Ok(Column::from_name(key.display_name()?)))
+ if let Some(col) = key.try_as_col() {
+ Ok(col.clone())
+ } else {
+ let name = key.display_name()?;
+ Ok(Column::from_name(name))
+ }
})
.collect::<Result<Vec<_>>>()?;
diff --git a/datafusion/expr/src/logical_plan/plan.rs
b/datafusion/expr/src/logical_plan/plan.rs
index 9832b69f84..266e7abc34 100644
--- a/datafusion/expr/src/logical_plan/plan.rs
+++ b/datafusion/expr/src/logical_plan/plan.rs
@@ -369,8 +369,18 @@ impl LogicalPlan {
// The join keys in using-join must be columns.
let columns =
on.iter().try_fold(HashSet::new(), |mut accumu, (l, r)| {
- accumu.insert(l.try_into_col()?);
- accumu.insert(r.try_into_col()?);
+ let Some(l) = l.try_as_col().cloned() else {
+ return internal_err!(
+ "Invalid join key. Expected column, found
{l:?}"
+ );
+ };
+ let Some(r) = r.try_as_col().cloned() else {
+ return internal_err!(
+ "Invalid join key. Expected column, found
{r:?}"
+ );
+ };
+ accumu.insert(l);
+ accumu.insert(r);
Result::<_, DataFusionError>::Ok(accumu)
})?;
using_columns.push(columns);
diff --git a/datafusion/optimizer/src/push_down_filter.rs
b/datafusion/optimizer/src/push_down_filter.rs
index 9ce135b0d6..57b38bd0d0 100644
--- a/datafusion/optimizer/src/push_down_filter.rs
+++ b/datafusion/optimizer/src/push_down_filter.rs
@@ -535,8 +535,8 @@ fn push_down_join(
.on
.iter()
.filter_map(|(l, r)| {
- let left_col = l.try_into_col().ok()?;
- let right_col = r.try_into_col().ok()?;
+ let left_col = l.try_as_col().cloned()?;
+ let right_col = r.try_as_col().cloned()?;
Some((left_col, right_col))
})
.collect::<Vec<_>>();
diff --git a/datafusion/optimizer/src/simplify_expressions/inlist_simplifier.rs
b/datafusion/optimizer/src/simplify_expressions/inlist_simplifier.rs
index 9dcb8ed155..c8638eb723 100644
--- a/datafusion/optimizer/src/simplify_expressions/inlist_simplifier.rs
+++ b/datafusion/optimizer/src/simplify_expressions/inlist_simplifier.rs
@@ -52,7 +52,7 @@ impl TreeNodeRewriter for ShortenInListSimplifier {
// expressions
list.len() == 1
|| list.len() <= THRESHOLD_INLINE_INLIST
- && expr.try_into_col().is_ok()
+ && expr.try_as_col().is_some()
)
{
let first_val = list[0].clone();
diff --git a/datafusion/proto/src/logical_plan/mod.rs
b/datafusion/proto/src/logical_plan/mod.rs
index a6352bcefc..83e58c3a22 100644
--- a/datafusion/proto/src/logical_plan/mod.rs
+++ b/datafusion/proto/src/logical_plan/mod.rs
@@ -45,8 +45,9 @@ use datafusion::{
prelude::SessionContext,
};
use datafusion_common::{
- context, internal_err, not_impl_err, parsers::CompressionTypeVariant,
- plan_datafusion_err, DataFusionError, Result, TableReference,
+ context, internal_datafusion_err, internal_err, not_impl_err,
+ parsers::CompressionTypeVariant, plan_datafusion_err, DataFusionError,
Result,
+ TableReference,
};
use datafusion_expr::{
dml,
@@ -695,7 +696,12 @@ impl AsLogicalPlan for LogicalPlanNode {
// The equijoin keys in using-join must be column.
let using_keys = left_keys
.into_iter()
- .map(|key| key.try_into_col())
+ .map(|key| {
+ key.try_as_col().cloned()
+ .ok_or_else(|| internal_datafusion_err!(
+ "Using join keys must be column
references, got: {key:?}"
+ ))
+ })
.collect::<Result<Vec<_>, _>>()?;
builder.join_using(
into_logical_plan!(join.right, ctx,
extension_codec)?,
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
