From 2d32cb27aa61c1a669f4b708fe5db223196d7af4 Mon Sep 17 00:00:00 2001 From: Dmitry Patsura Date: Mon, 16 May 2022 17:16:58 +0300 Subject: [PATCH] feat: Initial support for AnyExpression --- .../core/src/datasource/listing/helpers.rs | 1 + .../src/optimizer/common_subexpr_eliminate.rs | 4 + .../src/optimizer/simplify_expressions.rs | 1 + datafusion/core/src/optimizer/utils.rs | 8 + datafusion/core/src/physical_plan/planner.rs | 32 +- datafusion/core/tests/sql/expr.rs | 24 + datafusion/expr/src/expr.rs | 18 + datafusion/expr/src/expr_rewriter.rs | 5 + datafusion/expr/src/expr_schema.rs | 6 + datafusion/expr/src/expr_visitor.rs | 4 + datafusion/expr/src/utils.rs | 1 + .../physical-expr/src/expressions/any.rs | 674 ++++++++++++++++++ .../physical-expr/src/expressions/mod.rs | 2 + datafusion/proto/proto/datafusion.proto | 9 + datafusion/proto/src/from_proto.rs | 5 + datafusion/proto/src/to_proto.rs | 10 + datafusion/sql/src/planner.rs | 38 + datafusion/sql/src/utils.rs | 5 + 18 files changed, 844 insertions(+), 3 deletions(-) create mode 100644 datafusion/physical-expr/src/expressions/any.rs diff --git a/datafusion/core/src/datasource/listing/helpers.rs b/datafusion/core/src/datasource/listing/helpers.rs index a26eafabb50ca..317ba1d1a6949 100644 --- a/datafusion/core/src/datasource/listing/helpers.rs +++ b/datafusion/core/src/datasource/listing/helpers.rs @@ -86,6 +86,7 @@ impl ExpressionVisitor for ApplicabilityVisitor<'_> { | Expr::Cast { .. } | Expr::TryCast { .. } | Expr::BinaryExpr { .. } + | Expr::AnyExpr { .. } | Expr::Between { .. } | Expr::InList { .. } | Expr::Exists { .. } diff --git a/datafusion/core/src/optimizer/common_subexpr_eliminate.rs b/datafusion/core/src/optimizer/common_subexpr_eliminate.rs index 916e99713d4a6..ea1427363ac1d 100644 --- a/datafusion/core/src/optimizer/common_subexpr_eliminate.rs +++ b/datafusion/core/src/optimizer/common_subexpr_eliminate.rs @@ -404,6 +404,10 @@ impl ExprIdentifierVisitor<'_> { desc.push_str("BinaryExpr-"); desc.push_str(&op.to_string()); } + Expr::AnyExpr { op, .. } => { + desc.push_str("AnyExpr-"); + desc.push_str(&op.to_string()); + } Expr::Not(_) => { desc.push_str("Not-"); } diff --git a/datafusion/core/src/optimizer/simplify_expressions.rs b/datafusion/core/src/optimizer/simplify_expressions.rs index b85377df33fdc..a3c14e894a698 100644 --- a/datafusion/core/src/optimizer/simplify_expressions.rs +++ b/datafusion/core/src/optimizer/simplify_expressions.rs @@ -400,6 +400,7 @@ impl<'a> ConstEvaluator<'a> { Expr::ScalarUDF { fun, .. } => Self::volatility_ok(fun.signature.volatility), Expr::Literal(_) | Expr::BinaryExpr { .. } + | Expr::AnyExpr { .. } | Expr::Not(_) | Expr::IsNotNull(_) | Expr::IsNull(_) diff --git a/datafusion/core/src/optimizer/utils.rs b/datafusion/core/src/optimizer/utils.rs index 863536972f735..080d5f4129788 100644 --- a/datafusion/core/src/optimizer/utils.rs +++ b/datafusion/core/src/optimizer/utils.rs @@ -61,6 +61,9 @@ pub fn expr_sub_expressions(expr: &Expr) -> Result> { Expr::BinaryExpr { left, right, .. } => { Ok(vec![left.as_ref().to_owned(), right.as_ref().to_owned()]) } + Expr::AnyExpr { left, right, .. } => { + Ok(vec![left.as_ref().to_owned(), right.as_ref().to_owned()]) + } Expr::IsNull(expr) | Expr::IsNotNull(expr) | Expr::Cast { expr, .. } @@ -154,6 +157,11 @@ pub fn rewrite_expression(expr: &Expr, expressions: &[Expr]) -> Result { op: *op, right: Box::new(expressions[1].clone()), }), + Expr::AnyExpr { op, .. } => Ok(Expr::AnyExpr { + left: Box::new(expressions[0].clone()), + op: *op, + right: Box::new(expressions[1].clone()), + }), Expr::IsNull(_) => Ok(Expr::IsNull(Box::new(expressions[0].clone()))), Expr::IsNotNull(_) => Ok(Expr::IsNotNull(Box::new(expressions[0].clone()))), Expr::ScalarFunction { fun, .. } => Ok(Expr::ScalarFunction { diff --git a/datafusion/core/src/physical_plan/planner.rs b/datafusion/core/src/physical_plan/planner.rs index ad957409c8264..972ac003df15e 100644 --- a/datafusion/core/src/physical_plan/planner.rs +++ b/datafusion/core/src/physical_plan/planner.rs @@ -19,8 +19,15 @@ use super::analyze::AnalyzeExec; use super::{ - aggregates, empty::EmptyExec, expressions::binary, functions, - hash_join::PartitionMode, udaf, union::UnionExec, values::ValuesExec, windows, + aggregates, + empty::EmptyExec, + expressions::{any, binary}, + functions, + hash_join::PartitionMode, + udaf, + union::UnionExec, + values::ValuesExec, + windows, }; use crate::execution::context::{ExecutionProps, SessionState}; use crate::logical_expr::utils::generate_sort_key; @@ -109,6 +116,11 @@ fn create_physical_name(e: &Expr, is_first_expr: bool) -> Result { let right = create_physical_name(right, false)?; Ok(format!("{} {:?} {}", left, op, right)) } + Expr::AnyExpr { left, op, right } => { + let left = create_physical_name(left, false)?; + let right = create_physical_name(right, false)?; + Ok(format!("{} {:?} ANY({})", left, op, right)) + } Expr::Case { expr, when_then_expr, @@ -1182,7 +1194,6 @@ pub fn create_physical_expr( create_physical_expr(expr, input_dfschema, input_schema, execution_props)?, key.clone(), ))), - Expr::ScalarFunction { fun, args } => { let physical_args = args .iter() @@ -1245,6 +1256,21 @@ pub fn create_physical_expr( binary_expr } } + Expr::AnyExpr { left, op, right } => { + let lhs = create_physical_expr( + left, + input_dfschema, + input_schema, + execution_props, + )?; + let rhs = create_physical_expr( + right, + input_dfschema, + input_schema, + execution_props, + )?; + any(lhs, *op, rhs, input_schema) + } Expr::InList { expr, list, diff --git a/datafusion/core/tests/sql/expr.rs b/datafusion/core/tests/sql/expr.rs index 6046fc6d8b262..83e7b785f5c49 100644 --- a/datafusion/core/tests/sql/expr.rs +++ b/datafusion/core/tests/sql/expr.rs @@ -955,6 +955,30 @@ async fn test_extract_date_part() -> Result<()> { Ok(()) } +#[tokio::test] +async fn test_binary_any() -> Result<()> { + // = + test_expression!("1 = ANY([1, 2])", "true"); + test_expression!("3 = ANY([1, 2])", "false"); + test_expression!("NULL = ANY([1, 2])", "NULL"); + // utf8 + test_expression!("'a' = ANY(['a', 'b'])", "true"); + test_expression!("'c' = ANY(['a', 'b'])", "false"); + // bool + test_expression!("true = ANY([true, false])", "true"); + test_expression!("false = ANY([true, false])", "true"); + test_expression!("false = ANY([true, true])", "false"); + // <> + test_expression!("3 <> ANY([1, 2])", "true"); + test_expression!("1 <> ANY([1, 2])", "false"); + test_expression!("2 <> ANY([1, 2])", "false"); + test_expression!("NULL = ANY([1, 2])", "NULL"); + test_expression!("'c' <> ANY(['a', 'b'])", "true"); + test_expression!("'a' <> ANY(['a', 'b'])", "false"); + + Ok(()) +} + #[tokio::test] async fn test_in_list_scalar() -> Result<()> { test_expression!("'a' IN ('a','b')", "true"); diff --git a/datafusion/expr/src/expr.rs b/datafusion/expr/src/expr.rs index db6425a771481..60413cfd6435e 100644 --- a/datafusion/expr/src/expr.rs +++ b/datafusion/expr/src/expr.rs @@ -99,6 +99,15 @@ pub enum Expr { /// Right-hand side of the expression right: Box, }, + /// A binary expression such as "age > 21" + AnyExpr { + /// Left-hand side of the expression + left: Box, + /// The comparison operator + op: Operator, + /// Right-hand side of the expression + right: Box, + }, /// Negation of an expression. The expression's type must be a boolean to make sense. Not(Box), /// Whether an expression is not Null. This expression is never null. @@ -302,6 +311,7 @@ impl Expr { Expr::AggregateUDF { .. } => "AggregateUDF", Expr::Alias(..) => "Alias", Expr::Between { .. } => "Between", + Expr::AnyExpr { .. } => "AnyExpr", Expr::BinaryExpr { .. } => "BinaryExpr", Expr::Case { .. } => "Case", Expr::Cast { .. } => "Cast", @@ -524,6 +534,9 @@ impl fmt::Debug for Expr { Expr::BinaryExpr { left, op, right } => { write!(f, "{:?} {} {:?}", left, op, right) } + Expr::AnyExpr { left, op, right } => { + write!(f, "{:?} {} ANY({:?})", left, op, right) + } Expr::Sort { expr, asc, @@ -707,6 +720,11 @@ fn create_name(e: &Expr, input_schema: &DFSchema) -> Result { let right = create_name(right, input_schema)?; Ok(format!("{} {} {}", left, op, right)) } + Expr::AnyExpr { left, op, right } => { + let left = create_name(left, input_schema)?; + let right = create_name(right, input_schema)?; + Ok(format!("{} {} ANY({})", left, op, right)) + } Expr::Case { expr, when_then_expr, diff --git a/datafusion/expr/src/expr_rewriter.rs b/datafusion/expr/src/expr_rewriter.rs index e8cf049dde6f7..8ba76606eb67d 100644 --- a/datafusion/expr/src/expr_rewriter.rs +++ b/datafusion/expr/src/expr_rewriter.rs @@ -127,6 +127,11 @@ impl ExprRewritable for Expr { op, right: rewrite_boxed(right, rewriter)?, }, + Expr::AnyExpr { left, op, right } => Expr::AnyExpr { + left: rewrite_boxed(left, rewriter)?, + op, + right: rewrite_boxed(right, rewriter)?, + }, Expr::Not(expr) => Expr::Not(rewrite_boxed(expr, rewriter)?), Expr::IsNotNull(expr) => Expr::IsNotNull(rewrite_boxed(expr, rewriter)?), Expr::IsNull(expr) => Expr::IsNull(rewrite_boxed(expr, rewriter)?), diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index 2433024e38a4e..e110ce262a1a3 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -104,6 +104,7 @@ impl ExprSchemable for Expr { | Expr::InSubquery { .. } | Expr::Between { .. } | Expr::InList { .. } + | Expr::AnyExpr { .. } | Expr::IsNotNull(_) => Ok(DataType::Boolean), Expr::ScalarSubquery(subquery) => { Ok(subquery.subquery.schema().field(0).data_type().clone()) @@ -191,6 +192,11 @@ impl ExprSchemable for Expr { ref right, .. } => Ok(left.nullable(input_schema)? || right.nullable(input_schema)?), + Expr::AnyExpr { + ref left, + ref right, + .. + } => Ok(left.nullable(input_schema)? || right.nullable(input_schema)?), Expr::Wildcard => Err(DataFusionError::Internal( "Wildcard expressions are not valid in a logical query plan".to_owned(), )), diff --git a/datafusion/expr/src/expr_visitor.rs b/datafusion/expr/src/expr_visitor.rs index 05bc3bbfae537..8ff33f57a36b2 100644 --- a/datafusion/expr/src/expr_visitor.rs +++ b/datafusion/expr/src/expr_visitor.rs @@ -128,6 +128,10 @@ impl ExprVisitable for Expr { let visitor = left.accept(visitor)?; right.accept(visitor) } + Expr::AnyExpr { left, right, .. } => { + let visitor = left.accept(visitor)?; + right.accept(visitor) + } Expr::Between { expr, low, high, .. } => { diff --git a/datafusion/expr/src/utils.rs b/datafusion/expr/src/utils.rs index 483b12b49a5e8..5bcd28a90c3d2 100644 --- a/datafusion/expr/src/utils.rs +++ b/datafusion/expr/src/utils.rs @@ -58,6 +58,7 @@ impl ExpressionVisitor for ColumnNameVisitor<'_> { Expr::Alias(_, _) | Expr::Literal(_) | Expr::BinaryExpr { .. } + | Expr::AnyExpr { .. } | Expr::Not(_) | Expr::IsNotNull(_) | Expr::IsNull(_) diff --git a/datafusion/physical-expr/src/expressions/any.rs b/datafusion/physical-expr/src/expressions/any.rs new file mode 100644 index 0000000000000..dd04ccd075632 --- /dev/null +++ b/datafusion/physical-expr/src/expressions/any.rs @@ -0,0 +1,674 @@ +// 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. + +//! Any expression + +use std::any::Any; +use std::sync::Arc; + +use arrow::array::{ + BooleanArray, Int16Array, Int32Array, Int64Array, Int8Array, ListArray, + PrimitiveArray, UInt16Array, UInt32Array, UInt64Array, UInt8Array, +}; +use arrow::datatypes::ArrowPrimitiveType; +use arrow::{ + datatypes::{DataType, Schema}, + record_batch::RecordBatch, +}; + +use crate::expressions::try_cast; +use crate::PhysicalExpr; +use arrow::array::*; + +use datafusion_common::{DataFusionError, Result}; +use datafusion_expr::{ColumnarValue, Operator}; + +macro_rules! compare_op_scalar { + ($LEFT: expr, $LIST_VALUES:expr, $OP:expr, $LIST_VALUES_TYPE:ty, $LIST_FROM_SCALAR: expr) => {{ + let mut builder = BooleanBuilder::new($LEFT.len()); + + if $LIST_FROM_SCALAR { + for i in 0..$LEFT.len() { + if $LEFT.is_null(i) { + builder.append_null()?; + } else { + if $LIST_VALUES.is_null(0) { + builder.append_null()?; + } else { + builder.append_value($OP( + $LEFT.value(i), + $LIST_VALUES + .value(0) + .as_any() + .downcast_ref::<$LIST_VALUES_TYPE>() + .unwrap(), + ))?; + } + } + } + } else { + for i in 0..$LEFT.len() { + if $LEFT.is_null(i) { + builder.append_null()?; + } else { + if $LIST_VALUES.is_null(i) { + builder.append_null()?; + } else { + builder.append_value($OP( + $LEFT.value(i), + $LIST_VALUES + .value(i) + .as_any() + .downcast_ref::<$LIST_VALUES_TYPE>() + .unwrap(), + ))?; + } + } + } + } + + Ok(builder.finish()) + }}; +} + +macro_rules! make_primitive { + ($VALUES:expr, $IN_VALUES:expr, $NEGATED:expr, $TYPE:ident, $LIST_FROM_SCALAR: expr) => {{ + let left = $VALUES.as_any().downcast_ref::<$TYPE>().expect(&format!( + "Unable to downcast values to {}", + stringify!($TYPE) + )); + + if $NEGATED { + Ok(ColumnarValue::Array(Arc::new(neq_primitive( + left, + $IN_VALUES, + $LIST_FROM_SCALAR, + )?))) + } else { + Ok(ColumnarValue::Array(Arc::new(eq_primitive( + left, + $IN_VALUES, + $LIST_FROM_SCALAR, + )?))) + } + }}; +} + +fn eq_primitive( + array: &PrimitiveArray, + list: &ListArray, + list_from_scalar: bool, +) -> Result { + compare_op_scalar!( + array, + list, + |x, v: &PrimitiveArray| v.values().contains(&x), + PrimitiveArray, + list_from_scalar + ) +} + +fn neq_primitive( + array: &PrimitiveArray, + list: &ListArray, + list_from_scalar: bool, +) -> Result { + compare_op_scalar!( + array, + list, + |x, v: &PrimitiveArray| !v.values().contains(&x), + PrimitiveArray, + list_from_scalar + ) +} + +fn eq_bool( + array: &BooleanArray, + list: &ListArray, + list_from_scalar: bool, +) -> Result { + compare_op_scalar!( + array, + list, + |x, v: &BooleanArray| unsafe { + for i in 0..v.len() { + if v.value_unchecked(i) == x { + return true; + } + } + + false + }, + BooleanArray, + list_from_scalar + ) +} + +fn neq_bool( + array: &BooleanArray, + list: &ListArray, + list_from_scalar: bool, +) -> Result { + compare_op_scalar!( + array, + list, + |x, v: &BooleanArray| unsafe { + for i in 0..v.len() { + if v.value_unchecked(i) == x { + return false; + } + } + + true + }, + BooleanArray, + list_from_scalar + ) +} + +fn eq_utf8( + array: &GenericStringArray, + list: &ListArray, + list_from_scalar: bool, +) -> Result { + compare_op_scalar!( + array, + list, + |x, v: &GenericStringArray| unsafe { + for i in 0..v.len() { + if v.value_unchecked(i) == x { + return true; + } + } + + false + }, + GenericStringArray, + list_from_scalar + ) +} + +fn neq_utf8( + array: &GenericStringArray, + list: &ListArray, + list_from_scalar: bool, +) -> Result { + compare_op_scalar!( + array, + list, + |x, v: &GenericStringArray| unsafe { + for i in 0..v.len() { + if v.value_unchecked(i) == x { + return false; + } + } + + true + }, + GenericStringArray, + list_from_scalar + ) +} + +/// AnyExpr +#[derive(Debug)] +pub struct AnyExpr { + value: Arc, + op: Operator, + list: Arc, +} + +impl AnyExpr { + /// Create a new Any expression + pub fn new( + value: Arc, + op: Operator, + list: Arc, + ) -> Self { + Self { value, op, list } + } + + /// Compare for specific utf8 types + fn compare_utf8( + &self, + array: ArrayRef, + list: &ListArray, + negated: bool, + list_from_scalar: bool, + ) -> Result { + let array = array + .as_any() + .downcast_ref::>() + .unwrap(); + + if negated { + Ok(ColumnarValue::Array(Arc::new(neq_utf8( + array, + list, + list_from_scalar, + )?))) + } else { + Ok(ColumnarValue::Array(Arc::new(eq_utf8( + array, + list, + list_from_scalar, + )?))) + } + } + + /// Get the left side of the binary expression + pub fn left(&self) -> &Arc { + &self.value + } + + /// Get the right side of the binary expression + pub fn right(&self) -> &Arc { + &self.list + } + + /// Get the operator for this binary expression + pub fn op(&self) -> &Operator { + &self.op + } + + /// Compare for specific utf8 types + fn compare_bool( + &self, + array: ArrayRef, + list: &ListArray, + negated: bool, + list_from_scalar: bool, + ) -> Result { + let array = array.as_any().downcast_ref::().unwrap(); + + if negated { + Ok(ColumnarValue::Array(Arc::new(neq_bool( + array, + list, + list_from_scalar, + )?))) + } else { + Ok(ColumnarValue::Array(Arc::new(eq_bool( + array, + list, + list_from_scalar, + )?))) + } + } +} + +impl std::fmt::Display for AnyExpr { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "{} {} ANY({})", self.value, self.op, self.list) + } +} + +impl PhysicalExpr for AnyExpr { + /// Return a reference to Any that can be used for downcasting + fn as_any(&self) -> &dyn Any { + self + } + + fn data_type(&self, _input_schema: &Schema) -> Result { + Ok(DataType::Boolean) + } + + fn nullable(&self, input_schema: &Schema) -> Result { + self.value.nullable(input_schema) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + let value = match self.value.evaluate(batch)? { + ColumnarValue::Array(array) => array, + ColumnarValue::Scalar(scalar) => scalar.to_array(), + }; + + let (list, list_from_scalar) = match self.list.evaluate(batch)? { + ColumnarValue::Array(array) => (array, false), + ColumnarValue::Scalar(scalar) => (scalar.to_array(), true), + }; + let as_list = list + .as_any() + .downcast_ref::() + .expect("Unable to downcast list to ListArray"); + + let negated = match self.op { + Operator::Eq => false, + Operator::NotEq => true, + op => { + return Err(DataFusionError::NotImplemented(format!( + "Operator for ANY expression, actual: {:?}", + op + ))); + } + }; + + match value.data_type() { + DataType::Float16 => { + make_primitive!(value, as_list, negated, Float16Array, list_from_scalar) + } + DataType::Float32 => { + make_primitive!(value, as_list, negated, Float32Array, list_from_scalar) + } + DataType::Float64 => { + make_primitive!(value, as_list, negated, Float64Array, list_from_scalar) + } + DataType::Int8 => { + make_primitive!(value, as_list, negated, Int8Array, list_from_scalar) + } + DataType::Int16 => { + make_primitive!(value, as_list, negated, Int16Array, list_from_scalar) + } + DataType::Int32 => { + make_primitive!(value, as_list, negated, Int32Array, list_from_scalar) + } + DataType::Int64 => { + make_primitive!(value, as_list, negated, Int64Array, list_from_scalar) + } + DataType::UInt8 => { + make_primitive!(value, as_list, negated, UInt8Array, list_from_scalar) + } + DataType::UInt16 => { + make_primitive!(value, as_list, negated, UInt16Array, list_from_scalar) + } + DataType::UInt32 => { + make_primitive!(value, as_list, negated, UInt32Array, list_from_scalar) + } + DataType::UInt64 => { + make_primitive!(value, as_list, negated, UInt64Array, list_from_scalar) + } + DataType::Boolean => { + self.compare_bool(value, as_list, negated, list_from_scalar) + } + DataType::Utf8 => { + self.compare_utf8::(value, as_list, negated, list_from_scalar) + } + DataType::LargeUtf8 => { + self.compare_utf8::(value, as_list, negated, list_from_scalar) + } + datatype => Result::Err(DataFusionError::NotImplemented(format!( + "AnyExpr does not support datatype {:?}.", + datatype + ))), + } + } +} + +/// return two physical expressions that are optionally coerced to a +/// common type that the binary operator supports. +fn any_cast( + value: Arc, + _op: &Operator, + list: Arc, + input_schema: &Schema, +) -> Result<(Arc, Arc)> { + let list_dt = list.data_type(input_schema)?; + let list_type = match &list_dt { + DataType::List(f) => f.data_type(), + dt => return Err(DataFusionError::Execution(format!( + "Unexpected type on the right side of ANY expression. Must be a List, actual: {}", + dt + ))), + }; + + Ok((try_cast(value, input_schema, list_type.clone())?, list)) +} + +/// Creates an expression AnyExpr +pub fn any( + value: Arc, + op: Operator, + list: Arc, + input_schema: &Schema, +) -> Result> { + let (l, r) = any_cast(value, &op, list, input_schema)?; + Ok(Arc::new(AnyExpr::new(l, op, r))) +} + +#[cfg(test)] +mod tests { + use arrow::datatypes::Field; + + use super::*; + use crate::expressions::{col, lit}; + use datafusion_common::{Result, ScalarValue}; + + // applies the any expr to an input batch + macro_rules! execute_any { + ($BATCH:expr, $OP:expr, $EXPECTED:expr, $COL_A:expr, $COL_B:expr, $SCHEMA:expr) => {{ + let expr = any($COL_A, $OP, $COL_B, $SCHEMA).unwrap(); + let result = expr.evaluate(&$BATCH)?.into_array($BATCH.num_rows()); + let result = result + .as_any() + .downcast_ref::() + .expect("failed to downcast to BooleanArray"); + let expected = &BooleanArray::from($EXPECTED); + assert_eq!(expected, result); + }}; + } + + #[test] + fn any_int64_array_list() -> Result<()> { + let field_a = Field::new("a", DataType::Int64, true); + let field_b = Field::new( + "b", + DataType::List(Box::new(Field::new("item", DataType::Int64, true))), + true, + ); + + let schema = Schema::new(vec![field_a, field_b]); + let a = Int64Array::from(vec![Some(0), Some(3), None]); + let col_a = col("a", &schema)?; + + let values_builder = Int64Builder::new(3 * 3); + let mut builder = ListBuilder::new(values_builder); + + for _ in 0..3 { + builder.values().append_value(0).unwrap(); + builder.values().append_value(1).unwrap(); + builder.values().append_value(2).unwrap(); + builder.append(true).unwrap(); + } + + let b = builder.finish(); + let col_b = col("b", &schema)?; + + let batch = RecordBatch::try_new( + Arc::new(schema.clone()), + vec![Arc::new(a), Arc::new(b)], + )?; + + execute_any!( + batch, + Operator::Eq, + vec![Some(true), Some(false), None], + col_a.clone(), + col_b.clone(), + &schema + ); + + Ok(()) + } + + /// This tests emulates a simple type coercion, where the right side is the main type of expression + #[test] + fn any_int32_column_int64_array_list() -> Result<()> { + let field_a = Field::new("a", DataType::Int32, true); + let field_b = Field::new( + "b", + DataType::List(Box::new(Field::new("item", DataType::Int64, true))), + true, + ); + + let schema = Schema::new(vec![field_a, field_b]); + let a = Int32Array::from(vec![Some(0), Some(3), None]); + let col_a = col("a", &schema)?; + + let values_builder = Int64Builder::new(3 * 3); + let mut builder = ListBuilder::new(values_builder); + + for _ in 0..3 { + builder.values().append_value(0).unwrap(); + builder.values().append_value(1).unwrap(); + builder.values().append_value(2).unwrap(); + builder.append(true).unwrap(); + } + + let b = builder.finish(); + let col_b = col("b", &schema)?; + + let batch = RecordBatch::try_new( + Arc::new(schema.clone()), + vec![Arc::new(a), Arc::new(b)], + )?; + + execute_any!( + batch, + Operator::Eq, + vec![Some(true), Some(false), None], + col_a.clone(), + col_b.clone(), + &schema + ); + + Ok(()) + } + + // applies the any expr to an input batch and list (scalar) + macro_rules! execute_any_scalar { + ($BATCH:expr, $LIST:expr, $OP:expr, $EXPECTED:expr, $COL:expr, $SCHEMA:expr) => {{ + let expr = any($COL, $OP, $LIST, $SCHEMA).unwrap(); + let result = expr.evaluate(&$BATCH)?.into_array($BATCH.num_rows()); + let result = result + .as_any() + .downcast_ref::() + .expect("failed to downcast to BooleanArray"); + let expected = &BooleanArray::from($EXPECTED); + assert_eq!(expected, result); + }}; + } + + #[test] + fn any_int64_scalar_list() -> Result<()> { + let field_a = Field::new("a", DataType::Int64, true); + let schema = Schema::new(vec![field_a.clone()]); + let a = Int64Array::from(vec![Some(0), Some(3), None]); + let col_a = col("a", &schema)?; + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)])?; + + // expression: "a = ANY (0, 1, 2)" + let list = lit(ScalarValue::List( + Some(vec![ + ScalarValue::Int64(Some(0)), + ScalarValue::Int64(Some(1)), + ScalarValue::Int64(Some(2)), + ]), + Box::new(DataType::Int64), + )); + + let schema = &Schema::new(vec![ + field_a, + Field::new( + "b", + DataType::List(Box::new(Field::new("d", DataType::Int64, true))), + true, + ), + ]); + execute_any_scalar!( + batch, + list, + Operator::Eq, + vec![Some(true), Some(false), None], + col_a.clone(), + schema + ); + + Ok(()) + } + + #[test] + fn any_utf8_scalar_list() -> Result<()> { + let field_a = Field::new("a", DataType::Utf8, true); + let schema = Schema::new(vec![field_a.clone()]); + let a = StringArray::from(vec![Some("a"), Some("d"), None]); + let col_a = col("a", &schema)?; + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)])?; + + // expression: "a = ANY ('a', 'b', 'c')" + let list = lit(ScalarValue::List( + Some(vec![ + ScalarValue::Utf8(Some("a".to_string())), + ScalarValue::Utf8(Some("b".to_string())), + ScalarValue::Utf8(Some("c".to_string())), + ]), + Box::new(DataType::Utf8), + )); + + let schema = &Schema::new(vec![ + field_a, + Field::new( + "b", + DataType::List(Box::new(Field::new("d", DataType::Utf8, true))), + true, + ), + ]); + execute_any_scalar!( + batch, + list, + Operator::Eq, + vec![Some(true), Some(false), None], + col_a.clone(), + schema + ); + + Ok(()) + } + + #[test] + fn any_bool_scalar_list() -> Result<()> { + let field_a = Field::new("a", DataType::Boolean, true); + let schema = Schema::new(vec![field_a.clone()]); + let a = BooleanArray::from(vec![Some(true), Some(false), None]); + let col_a = col("a", &schema)?; + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)])?; + + // expression: "a = ANY (true)" + let list = lit(ScalarValue::List( + Some(vec![ScalarValue::Boolean(Some(true))]), + Box::new(DataType::Boolean), + )); + + let schema = &Schema::new(vec![ + field_a, + Field::new( + "b", + DataType::List(Box::new(Field::new("d", DataType::Boolean, true))), + true, + ), + ]); + execute_any_scalar!( + batch, + list, + Operator::Eq, + vec![Some(true), Some(false), None], + col_a.clone(), + schema + ); + + Ok(()) + } +} diff --git a/datafusion/physical-expr/src/expressions/mod.rs b/datafusion/physical-expr/src/expressions/mod.rs index 22e80e20ec9fc..623b7bfe1d117 100644 --- a/datafusion/physical-expr/src/expressions/mod.rs +++ b/datafusion/physical-expr/src/expressions/mod.rs @@ -19,6 +19,7 @@ #[macro_use] mod binary; +mod any; mod case; mod cast; mod column; @@ -65,6 +66,7 @@ pub use crate::window::nth_value::NthValue; pub use crate::window::rank::{dense_rank, percent_rank, rank}; pub use crate::window::row_number::RowNumber; +pub use any::{any, AnyExpr}; pub use binary::{binary, BinaryExpr}; pub use case::{case, CaseExpr}; pub use cast::{ diff --git a/datafusion/proto/proto/datafusion.proto b/datafusion/proto/proto/datafusion.proto index d578714422b84..8d5ac643de697 100644 --- a/datafusion/proto/proto/datafusion.proto +++ b/datafusion/proto/proto/datafusion.proto @@ -301,6 +301,9 @@ message LogicalExprNode { // Scalar UDF expressions ScalarUDFExprNode scalar_udf_expr = 20; + + // binary expressions + AnyExprNode any_expr = 21; } } @@ -327,6 +330,12 @@ message BinaryExprNode { string op = 3; } +message AnyExprNode { + LogicalExprNode l = 1; + LogicalExprNode r = 2; + string op = 3; +} + message NegativeNode { LogicalExprNode expr = 1; } diff --git a/datafusion/proto/src/from_proto.rs b/datafusion/proto/src/from_proto.rs index 6717e5b47a5ac..39eb1e3adf453 100644 --- a/datafusion/proto/src/from_proto.rs +++ b/datafusion/proto/src/from_proto.rs @@ -928,6 +928,11 @@ pub fn parse_expr( op: from_proto_binary_op(&binary_expr.op)?, right: Box::new(parse_required_expr(&binary_expr.r, registry, "r")?), }), + ExprType::AnyExpr(any_expr) => Ok(Expr::AnyExpr { + left: Box::new(parse_required_expr(&any_expr.l, registry, "l")?), + op: from_proto_binary_op(&any_expr.op)?, + right: Box::new(parse_required_expr(&any_expr.r, registry, "r")?), + }), ExprType::Column(column) => Ok(Expr::Column(column.into())), ExprType::Literal(literal) => { let scalar_value: ScalarValue = literal.try_into()?; diff --git a/datafusion/proto/src/to_proto.rs b/datafusion/proto/src/to_proto.rs index 70c4d57c4c1d7..a24f5ecdd6c24 100644 --- a/datafusion/proto/src/to_proto.rs +++ b/datafusion/proto/src/to_proto.rs @@ -448,6 +448,16 @@ impl TryFrom<&Expr> for protobuf::LogicalExprNode { expr_type: Some(ExprType::BinaryExpr(binary_expr)), } } + Expr::AnyExpr { left, op, right } => { + let binary_expr = Box::new(protobuf::AnyExprNode { + l: Some(Box::new(left.as_ref().try_into()?)), + r: Some(Box::new(right.as_ref().try_into()?)), + op: format!("{:?}", op), + }); + Self { + expr_type: Some(ExprType::AnyExpr(binary_expr)), + } + } Expr::WindowFunction { ref fun, ref args, diff --git a/datafusion/sql/src/planner.rs b/datafusion/sql/src/planner.rs index 238cf7dbeb7a8..cbb556536b9e5 100644 --- a/datafusion/sql/src/planner.rs +++ b/datafusion/sql/src/planner.rs @@ -1550,6 +1550,30 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { } } + fn parse_sql_binary_any( + &self, + left: SQLExpr, + op: BinaryOperator, + right: SQLExpr, + schema: &DFSchema, + ctes: &mut HashMap, + ) -> Result { + let operator = match op { + BinaryOperator::Eq => Ok(Operator::Eq), + BinaryOperator::NotEq => Ok(Operator::NotEq), + _ => Err(DataFusionError::NotImplemented(format!( + "Unsupported SQL ANY operator {:?}", + op + ))), + }?; + + Ok(Expr::AnyExpr { + left: Box::new(self.sql_expr_to_logical_expr(left, schema, ctes)?), + op: operator, + right: Box::new(self.sql_expr_to_logical_expr(right, schema, ctes)?), + }) + } + fn parse_sql_binary_op( &self, left: SQLExpr, @@ -1558,6 +1582,20 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { schema: &DFSchema, ctes: &mut HashMap, ) -> Result { + match right { + SQLExpr::AnyOp(any_expr) => { + return self.parse_sql_binary_any(left, op, *any_expr, schema, ctes); + } + // TODO: https://github.com/apache/arrow-datafusion/issues/2547 + SQLExpr::AllOp(_) => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported SQL ALL operator {:?}", + right + ))); + } + _ => {} + }; + let operator = match op { BinaryOperator::Gt => Ok(Operator::Gt), BinaryOperator::GtEq => Ok(Operator::GtEq), diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index 92e98813501b4..f2cf6cbf29089 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -232,6 +232,11 @@ where op: *op, right: Box::new(clone_with_replacement(&**right, replacement_fn)?), }), + Expr::AnyExpr { left, right, op } => Ok(Expr::AnyExpr { + left: Box::new(clone_with_replacement(&**left, replacement_fn)?), + op: *op, + right: Box::new(clone_with_replacement(&**right, replacement_fn)?), + }), Expr::Case { expr: case_expr_opt, when_then_expr,