Skip to content

Commit b3ed4a8

Browse files
ovrovr
committed
feat: Initial support for AnyExpression
1 parent 934af5c commit b3ed4a8

File tree

14 files changed

+751
-2
lines changed

14 files changed

+751
-2
lines changed

‎datafusion/core/src/datasource/listing/helpers.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,7 @@ impl ExpressionVisitor for ApplicabilityVisitor<'_> {
9191
| Expr::Cast { .. }
9292
| Expr::TryCast { .. }
9393
| Expr::BinaryExpr { .. }
94+
| Expr::AnyExpr { .. }
9495
| Expr::Between { .. }
9596
| Expr::InList { .. }
9697
| Expr::GetIndexedField { .. }

‎datafusion/core/src/logical_plan/expr_rewriter.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,11 @@ impl ExprRewritable for Expr {
119119
op,
120120
right: rewrite_boxed(right, rewriter)?,
121121
},
122+
Expr::AnyExpr { left, op, right } => Expr::AnyExpr {
123+
left: rewrite_boxed(left, rewriter)?,
124+
op,
125+
right: rewrite_boxed(right, rewriter)?,
126+
},
122127
Expr::Not(expr) => Expr::Not(rewrite_boxed(expr, rewriter)?),
123128
Expr::IsNotNull(expr) => Expr::IsNotNull(rewrite_boxed(expr, rewriter)?),
124129
Expr::IsNull(expr) => Expr::IsNull(rewrite_boxed(expr, rewriter)?),

‎datafusion/core/src/logical_plan/expr_schema.rs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,7 @@ impl ExprSchemable for Expr {
111111
| Expr::IsNull(_)
112112
| Expr::Between { .. }
113113
| Expr::InList { .. }
114+
| Expr::AnyExpr { .. }
114115
| Expr::IsNotNull(_) => Ok(DataType::Boolean),
115116
Expr::BinaryExpr {
116117
ref left,
@@ -189,6 +190,11 @@ impl ExprSchemable for Expr {
189190
ref right,
190191
..
191192
} => Ok(left.nullable(input_schema)? || right.nullable(input_schema)?),
193+
Expr::AnyExpr {
194+
ref left,
195+
ref right,
196+
..
197+
} => Ok(left.nullable(input_schema)? || right.nullable(input_schema)?),
192198
Expr::Wildcard => Err(DataFusionError::Internal(
193199
"Wildcard expressions are not valid in a logical query plan".to_owned(),
194200
)),

‎datafusion/core/src/logical_plan/expr_visitor.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,10 @@ impl ExprVisitable for Expr {
116116
let visitor = left.accept(visitor)?;
117117
right.accept(visitor)
118118
}
119+
Expr::AnyExpr { left, right, .. } => {
120+
let visitor = left.accept(visitor)?;
121+
right.accept(visitor)
122+
}
119123
Expr::Between {
120124
expr, low, high, ..
121125
} => {

‎datafusion/core/src/optimizer/common_subexpr_eliminate.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -423,6 +423,10 @@ impl ExprIdentifierVisitor<'_> {
423423
desc.push_str("BinaryExpr-");
424424
desc.push_str(&op.to_string());
425425
}
426+
Expr::AnyExpr { op, .. } => {
427+
desc.push_str("AnyExpr-");
428+
desc.push_str(&op.to_string());
429+
}
426430
Expr::Not(_) => {
427431
desc.push_str("Not-");
428432
}

‎datafusion/core/src/optimizer/simplify_expressions.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -385,6 +385,7 @@ impl<'a> ConstEvaluator<'a> {
385385
Expr::TableUDF { .. } => false,
386386
Expr::Literal(_)
387387
| Expr::BinaryExpr { .. }
388+
| Expr::AnyExpr { .. }
388389
| Expr::Not(_)
389390
| Expr::IsNotNull(_)
390391
| Expr::IsNull(_)

‎datafusion/core/src/optimizer/utils.rs‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,7 @@ impl ExpressionVisitor for ColumnNameVisitor<'_> {
7373
Expr::Alias(_, _)
7474
| Expr::Literal(_)
7575
| Expr::BinaryExpr { .. }
76+
| Expr::AnyExpr { .. }
7677
| Expr::Not(_)
7778
| Expr::IsNotNull(_)
7879
| Expr::IsNull(_)
@@ -305,6 +306,9 @@ pub fn expr_sub_expressions(expr: &Expr) -> Result<Vec<Expr>> {
305306
Expr::BinaryExpr { left, right, .. } => {
306307
Ok(vec![left.as_ref().to_owned(), right.as_ref().to_owned()])
307308
}
309+
Expr::AnyExpr { left, right, .. } => {
310+
Ok(vec![left.as_ref().to_owned(), right.as_ref().to_owned()])
311+
}
308312
Expr::IsNull(expr)
309313
| Expr::IsNotNull(expr)
310314
| Expr::Cast { expr, .. }
@@ -394,6 +398,11 @@ pub fn rewrite_expression(expr: &Expr, expressions: &[Expr]) -> Result<Expr> {
394398
op: *op,
395399
right: Box::new(expressions[1].clone()),
396400
}),
401+
Expr::AnyExpr { op, .. } => Ok(Expr::AnyExpr {
402+
left: Box::new(expressions[0].clone()),
403+
op: *op,
404+
right: Box::new(expressions[1].clone()),
405+
}),
397406
Expr::IsNull(_) => Ok(Expr::IsNull(Box::new(expressions[0].clone()))),
398407
Expr::IsNotNull(_) => Ok(Expr::IsNotNull(Box::new(expressions[0].clone()))),
399408
Expr::ScalarFunction { fun, .. } => Ok(Expr::ScalarFunction {

‎datafusion/core/src/physical_plan/planner.rs‎

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,7 @@ use arrow::datatypes::{Schema, SchemaRef};
6767
use arrow::{compute::can_cast_types, datatypes::DataType};
6868
use async_trait::async_trait;
6969
use datafusion_common::OuterQueryCursor;
70-
use datafusion_physical_expr::expressions::OuterColumn;
70+
use datafusion_physical_expr::expressions::{any, OuterColumn};
7171
use futures::future::BoxFuture;
7272
use futures::{FutureExt, StreamExt, TryStreamExt};
7373
use log::{debug, trace};
@@ -112,6 +112,11 @@ fn create_physical_name(e: &Expr, is_first_expr: bool) -> Result<String> {
112112
let right = create_physical_name(right, false)?;
113113
Ok(format!("{} {:?} {}", left, op, right))
114114
}
115+
Expr::AnyExpr { left, op, right } => {
116+
let left = create_physical_name(left, false)?;
117+
let right = create_physical_name(right, false)?;
118+
Ok(format!("{} {:?} ANY({})", left, op, right))
119+
}
115120
Expr::Case {
116121
expr,
117122
when_then_expr,
@@ -1096,7 +1101,6 @@ pub fn create_physical_expr(
10961101
create_physical_expr(expr, input_dfschema, input_schema, execution_props)?,
10971102
create_physical_expr(key, input_dfschema, input_schema, execution_props)?,
10981103
))),
1099-
11001104
Expr::ScalarFunction { fun, args } => {
11011105
let physical_args = args
11021106
.iter()
@@ -1172,6 +1176,21 @@ pub fn create_physical_expr(
11721176
binary_expr
11731177
}
11741178
}
1179+
Expr::AnyExpr { left, op, right } => {
1180+
let lhs = create_physical_expr(
1181+
left,
1182+
input_dfschema,
1183+
input_schema,
1184+
execution_props,
1185+
)?;
1186+
let rhs = create_physical_expr(
1187+
right,
1188+
input_dfschema,
1189+
input_schema,
1190+
execution_props,
1191+
)?;
1192+
any(lhs, *op, rhs, input_schema)
1193+
}
11751194
Expr::InList {
11761195
expr,
11771196
list,

‎datafusion/core/src/sql/planner.rs‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1477,13 +1477,49 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
14771477
}
14781478
}
14791479

1480+
fn parse_sql_binary_any(
1481+
&self,
1482+
left: SQLExpr,
1483+
op: BinaryOperator,
1484+
right: Box<SQLExpr>,
1485+
schema: &DFSchema,
1486+
) -> Result<Expr> {
1487+
let operator = match op {
1488+
BinaryOperator::Eq => Ok(Operator::Eq),
1489+
BinaryOperator::NotEq => Ok(Operator::NotEq),
1490+
_ => Err(DataFusionError::NotImplemented(format!(
1491+
"Unsupported SQL ANY operator {:?}",
1492+
op
1493+
))),
1494+
}?;
1495+
1496+
Ok(Expr::AnyExpr {
1497+
left: Box::new(self.sql_expr_to_logical_expr(left, schema)?),
1498+
op: operator,
1499+
right: Box::new(self.sql_expr_to_logical_expr(*right, schema)?),
1500+
})
1501+
}
1502+
14801503
fn parse_sql_binary_op(
14811504
&self,
14821505
left: SQLExpr,
14831506
op: BinaryOperator,
14841507
right: SQLExpr,
14851508
schema: &DFSchema,
14861509
) -> Result<Expr> {
1510+
match right {
1511+
SQLExpr::AnyOp(any_expr) => {
1512+
return self.parse_sql_binary_any(left, op, any_expr, schema);
1513+
}
1514+
SQLExpr::AllOp(_) => {
1515+
return Err(DataFusionError::NotImplemented(format!(
1516+
"Unsupported SQL ALL operator {:?}",
1517+
right
1518+
)));
1519+
}
1520+
_ => {}
1521+
};
1522+
14871523
let operator = match op {
14881524
BinaryOperator::Gt => Ok(Operator::Gt),
14891525
BinaryOperator::GtEq => Ok(Operator::GtEq),

‎datafusion/core/src/sql/utils.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -297,6 +297,11 @@ where
297297
op: *op,
298298
right: Box::new(clone_with_replacement(&**right, replacement_fn)?),
299299
}),
300+
Expr::AnyExpr { left, right, op } => Ok(Expr::AnyExpr {
301+
left: Box::new(clone_with_replacement(&**left, replacement_fn)?),
302+
op: *op,
303+
right: Box::new(clone_with_replacement(&**right, replacement_fn)?),
304+
}),
300305
Expr::Case {
301306
expr: case_expr_opt,
302307
when_then_expr,

0 commit comments

Comments
 (0)