Merge pull request #7712 from sundy-li/Kleene-logic-and

mergify[bot] · web-flow · commit 0175e2002b91 · 2022-09-20T03:29:26.000Z
feat(query): add and_filters function
diff --git a/src/query/expression/src/types/nullable.rs b/src/query/expression/src/types/nullable.rs
@@ -69,6 +69,13 @@ impl<T: ValueType> ValueType for NullableType<T> {
                 has_null: *has_null,
                 value: Some(Box::new(T::try_downcast_domain(value)?)),
             }),
+            Domain::Nullable(NullableDomain {
+                has_null,
+                value: None,
+            }) => Some(NullableDomain {
+                has_null: *has_null,
+                value: None,
+            }),
             _ => None,
         }
     }
diff --git a/src/query/functions-v2/src/scalars/boolean.rs b/src/query/functions-v2/src/scalars/boolean.rs
@@ -13,9 +13,14 @@
 // limitations under the License.
 
 use common_expression::types::boolean::BooleanDomain;
+use common_expression::types::nullable::NullableDomain;
 use common_expression::types::BooleanType;
+use common_expression::types::NullableType;
+use common_expression::vectorize_2_arg;
 use common_expression::FunctionProperty;
 use common_expression::FunctionRegistry;
+use common_expression::Value;
+use common_expression::ValueRef;
 
 pub fn register(registry: &mut FunctionRegistry) {
     registry.register_1_arg::<BooleanType, BooleanType, _, _>(
@@ -29,7 +34,37 @@ pub fn register(registry: &mut FunctionRegistry) {
         },
         |val| !val,
     );
-    registry.register_2_arg::<BooleanType, BooleanType, BooleanType, _, _>(
+
+    // special function to combine the filter efficiently
+    registry.register_passthrough_nullable_2_arg::<BooleanType, BooleanType, BooleanType, _, _>(
+        "and_filters",
+        FunctionProperty::default(),
+        |lhs, rhs| {
+            Some(BooleanDomain {
+                has_false: lhs.has_false || rhs.has_false,
+                has_true: lhs.has_true && rhs.has_true,
+            })
+        },
+        |lhs, rhs, _| match (lhs, rhs) {
+            (ValueRef::Scalar(flag), other) => {
+                if flag {
+                    Ok(other.to_owned())
+                } else {
+                    Ok(Value::Scalar(false))
+                }
+            }
+            (other, ValueRef::Scalar(flag)) => {
+                if flag {
+                    Ok(other.to_owned())
+                } else {
+                    Ok(Value::Scalar(false))
+                }
+            }
+            (ValueRef::Column(a), ValueRef::Column(b)) => Ok(Value::Column(&a & &b)),
+        },
+    );
+
+    registry.register_2_arg_core::<BooleanType, BooleanType, BooleanType, _, _>(
         "and",
         FunctionProperty::default(),
         |lhs, rhs| {
@@ -38,9 +73,26 @@ pub fn register(registry: &mut FunctionRegistry) {
                 has_true: lhs.has_true && rhs.has_true,
             })
         },
-        |lhs, rhs| lhs && rhs,
+        |lhs, rhs, _| match (lhs, rhs) {
+            (ValueRef::Scalar(flag), other) => {
+                if flag {
+                    Ok(other.to_owned())
+                } else {
+                    Ok(Value::Scalar(false))
+                }
+            }
+            (other, ValueRef::Scalar(flag)) => {
+                if flag {
+                    Ok(other.to_owned())
+                } else {
+                    Ok(Value::Scalar(false))
+                }
+            }
+            (ValueRef::Column(a), ValueRef::Column(b)) => Ok(Value::Column(&a & &b)),
+        },
     );
-    registry.register_2_arg::<BooleanType, BooleanType, BooleanType, _, _>(
+
+    registry.register_2_arg_core::<BooleanType, BooleanType, BooleanType, _, _>(
         "or",
         FunctionProperty::default(),
         |lhs, rhs| {
@@ -49,8 +101,88 @@ pub fn register(registry: &mut FunctionRegistry) {
                 has_true: lhs.has_true || rhs.has_true,
             })
         },
-        |lhs, rhs| lhs || rhs,
+        |lhs, rhs, _| match (lhs, rhs) {
+            (ValueRef::Scalar(flag), other) => {
+                if flag {
+                    Ok(Value::Scalar(true))
+                } else {
+                    Ok(other.to_owned())
+                }
+            }
+            (other, ValueRef::Scalar(flag)) => {
+                if flag {
+                    Ok(Value::Scalar(true))
+                } else {
+                    Ok(other.to_owned())
+                }
+            }
+            (ValueRef::Column(a), ValueRef::Column(b)) => Ok(Value::Column(&a | &b)),
+        },
+    );
+
+    // https://en.wikibooks.org/wiki/Structured_Query_Language/NULLs_and_the_Three_Valued_Logic
+    registry.register_2_arg_core::<NullableType<BooleanType>, NullableType<BooleanType>, NullableType<BooleanType>, _, _>(
+        "and",
+        FunctionProperty::default(),
+        |lhs, rhs| {
+            if !lhs.has_null && !rhs.has_null {
+                let bools = match (&lhs.value, &rhs.value) {
+                    (Some(a), Some(b)) => Some(Box::new(BooleanDomain {
+                    has_false: a.has_false || b.has_false,
+                    has_true: a.has_true && b.has_true,
+                    })),
+                    _ => return None,
+                };
+                return Some(NullableDomain::<BooleanType> {
+                    has_null: false,
+                    value: bools,
+                });
+            }
+            None
+        },
+        // value = lhs & rhs,  valid = (lhs_v & rhs_v) | (!lhs & lhs_v) | (!rhs & rhs_v))
+        vectorize_2_arg::<NullableType<BooleanType>, NullableType<BooleanType>, NullableType<BooleanType>>(|lhs, rhs| {
+            let lhs_v = lhs.is_some();
+            let rhs_v = rhs.is_some();
+            let valid = (lhs_v & rhs_v) | (lhs == Some(false)) | (rhs == Some(false));
+            if valid {
+                Some(lhs.unwrap_or_default() & rhs.unwrap_or_default())
+            } else {
+                None
+            }
+        }),
     );
+
+    registry.register_2_arg_core::<NullableType<BooleanType>, NullableType<BooleanType>, NullableType<BooleanType>, _, _>(
+        "or",
+        FunctionProperty::default(),
+        |lhs, rhs| {
+            if !lhs.has_null && !rhs.has_null {
+                let bools = match (&lhs.value, &rhs.value) {
+                    (Some(a), Some(b)) => Some(Box::new(BooleanDomain {
+                        has_false: a.has_false && b.has_false,
+                        has_true: a.has_true || b.has_true,
+                    })),
+                    _ => return None,
+                };
+                return Some(NullableDomain::<BooleanType> {
+                    has_null: false,
+                    value: bools,
+                });
+            }
+            None
+        },
+        // value = lhs | rhs,  valid = (lhs_v & rhs_v) | (lhs | rhs))
+        vectorize_2_arg::<NullableType<BooleanType>, NullableType<BooleanType>, NullableType<BooleanType>>(|lhs, rhs| {
+            let valid = (lhs.is_some() & rhs.is_some()) | (lhs.unwrap_or_default() | rhs.unwrap_or_default());
+            if valid {
+                Some(lhs.unwrap_or_default() | rhs.unwrap_or_default())
+            } else {
+                None
+            }
+        }),
+    );
+
     registry.register_2_arg::<BooleanType, BooleanType, BooleanType, _, _>(
         "xor",
         FunctionProperty::default(),
diff --git a/src/query/functions-v2/tests/it/scalars/boolean.rs b/src/query/functions-v2/tests/it/scalars/boolean.rs
@@ -34,7 +34,11 @@ fn test_boolean() {
 
 fn test_and(file: &mut impl Write) {
     run_ast(file, "true AND false", &[]);
-    run_ast(file, "null AND false", &[]);
+    run_ast(file, "true AND null", &[]);
+    run_ast(file, "true AND true", &[]);
+    run_ast(file, "false AND false", &[]);
+    run_ast(file, "false AND null", &[]);
+    run_ast(file, "false AND true", &[]);
 }
 
 fn test_not(file: &mut impl Write) {
diff --git a/src/query/functions-v2/tests/it/scalars/testdata/boolean.txt b/src/query/functions-v2/tests/it/scalars/testdata/boolean.txt
@@ -7,15 +7,51 @@ output domain  : {FALSE}
 output         : false
 
 
-ast            : null AND false
-raw expr       : and(NULL, false)
-checked expr   : and<NULL, Boolean NULL>(NULL, CAST(false AS Boolean NULL))
+ast            : true AND null
+raw expr       : and(true, NULL)
+checked expr   : and<Boolean NULL, Boolean NULL>(CAST(true AS Boolean NULL), CAST(NULL AS Boolean NULL))
 optimized expr : NULL
-output type    : NULL
-output domain  : {NULL}
+output type    : Boolean NULL
+output domain  : Unknown
 output         : NULL
 
 
+ast            : true AND true
+raw expr       : and(true, true)
+checked expr   : and<Boolean, Boolean>(true, true)
+optimized expr : true
+output type    : Boolean
+output domain  : {TRUE}
+output         : true
+
+
+ast            : false AND false
+raw expr       : and(false, false)
+checked expr   : and<Boolean, Boolean>(false, false)
+optimized expr : false
+output type    : Boolean
+output domain  : {FALSE}
+output         : false
+
+
+ast            : false AND null
+raw expr       : and(false, NULL)
+checked expr   : and<Boolean NULL, Boolean NULL>(CAST(false AS Boolean NULL), CAST(NULL AS Boolean NULL))
+optimized expr : false
+output type    : Boolean NULL
+output domain  : Unknown
+output         : false
+
+
+ast            : false AND true
+raw expr       : and(false, true)
+checked expr   : and<Boolean, Boolean>(false, true)
+optimized expr : false
+output type    : Boolean
+output domain  : {FALSE}
+output         : false
+
+
 ast            : NOT a
 raw expr       : not(ColumnRef(0)::NULL)
 checked expr   : not<NULL>(ColumnRef(0))
@@ -102,10 +138,10 @@ output         : true
 
 ast            : null OR false
 raw expr       : or(NULL, false)
-checked expr   : or<NULL, Boolean NULL>(NULL, CAST(false AS Boolean NULL))
+checked expr   : or<Boolean NULL, Boolean NULL>(CAST(NULL AS Boolean NULL), CAST(false AS Boolean NULL))
 optimized expr : NULL
-output type    : NULL
-output domain  : {NULL}
+output type    : Boolean NULL
+output domain  : Unknown
 output         : NULL
 
 
diff --git a/src/query/functions/src/scalars/logics/and.rs b/src/query/functions/src/scalars/logics/and.rs
@@ -12,7 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use core::fmt;
+
+use common_arrow::arrow::bitmap::MutableBitmap;
 use common_datavalues::prelude::*;
+use common_exception::ErrorCode;
 use common_exception::Result;
 
 use super::logic::LogicExpression;
@@ -30,6 +34,14 @@ impl_logic_expression!(LogicAndExpression, &, |lhs: bool, rhs: bool, lhs_v: bool
     (lhs & rhs,  (lhs_v & rhs_v) | (!lhs & lhs_v) | (!rhs & rhs_v))
 });
 
+/// Logical functions AND, OR, XOR and NOT support three-valued (or ternary) logic
+/// https://en.wikibooks.org/wiki/Structured_Query_Language/NULLs_and_the_Three_Valued_Logic
+///
+/// Functions XOR and NOT rely on "default implementation for NULLs":
+///   - if any of the arguments is of Nullable type, the return value type is Nullable
+///   - if any of the arguments is NULL, the return value is NULL
+///
+/// Functions AND and OR provide their own special implementations for ternary logic
 #[derive(Clone)]
 pub struct LogicAndFunction;
 
@@ -47,3 +59,79 @@ impl LogicAndFunction {
         )
     }
 }
+
+#[derive(Clone)]
+pub struct LogicAndFiltersFunction;
+
+impl Function for LogicAndFiltersFunction {
+    fn name(&self) -> &str {
+        "and_filters"
+    }
+
+    fn return_type(&self) -> DataTypeImpl {
+        bool::to_data_type()
+    }
+
+    fn eval(
+        &self,
+        _func_ctx: FunctionContext,
+        columns: &ColumnsWithField,
+        input_rows: usize,
+    ) -> Result<ColumnRef> {
+        if columns.len() == 1 {
+            return Ok(columns[1].column().clone());
+        }
+
+        let mut validity = None;
+        for c in columns.iter() {
+            let c = c.column();
+            if c.is_const() {
+                let v = c.get_bool(0)?;
+                if !v {
+                    let validity = MutableBitmap::from_len_zeroed(input_rows).into();
+                    return Ok(BooleanColumn::from_arrow_data(validity).arc());
+                }
+            } else {
+                let bools: &BooleanColumn = Series::check_get(c)?;
+                match validity.as_mut() {
+                    Some(v) => {
+                        *v = &*v & bools.values();
+                    }
+                    None => validity = Some(bools.values().clone()),
+                }
+            }
+        }
+        let validity = validity.unwrap_or_else(|| MutableBitmap::from_len_set(input_rows).into());
+        Ok(BooleanColumn::from_arrow_data(validity).arc())
+    }
+}
+
+impl LogicAndFiltersFunction {
+    pub fn try_create(_display_name: &str, args: &[&DataTypeImpl]) -> Result<Box<dyn Function>> {
+        if args
+            .iter()
+            .any(|arg| !matches!(arg, DataTypeImpl::Boolean(_)))
+        {
+            return Err(ErrorCode::IllegalDataType(format!(
+                "Illegal type {:?} of argument of function and_filters, expect to be all boolean types",
+                args
+            )));
+        }
+
+        Ok(Box::new(Self))
+    }
+
+    pub fn desc() -> FunctionDescription {
+        FunctionDescription::creator(Box::new(Self::try_create)).features(
+            FunctionFeatures::default()
+                .deterministic()
+                .variadic_arguments(1, std::usize::MAX),
+        )
+    }
+}
+
+impl fmt::Display for LogicAndFiltersFunction {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "and_filters")
+    }
+}
diff --git a/src/query/functions/src/scalars/logics/logic.rs b/src/query/functions/src/scalars/logics/logic.rs
@@ -18,6 +18,7 @@ use common_datavalues::prelude::*;
 use common_exception::Result;
 
 use super::xor::LogicXorFunction;
+use super::LogicAndFiltersFunction;
 use super::LogicAndFunction;
 use super::LogicNotFunction;
 use super::LogicOrFunction;
@@ -31,6 +32,7 @@ pub struct LogicFunction;
 impl LogicFunction {
     pub fn register(factory: &mut FunctionFactory) {
         factory.register("and", LogicAndFunction::desc());
+        factory.register("and_filters", LogicAndFiltersFunction::desc());
         factory.register("or", LogicOrFunction::desc());
         factory.register("not", LogicNotFunction::desc());
         factory.register("xor", LogicXorFunction::desc());
diff --git a/src/query/functions/src/scalars/logics/mod.rs b/src/query/functions/src/scalars/logics/mod.rs
diff --git a/src/query/service/src/sql/executor/pipeline_builder.rs b/src/query/service/src/sql/executor/pipeline_builder.rs

Original file line number	Diff line number	Diff line change
`@@ -69,6 +69,13 @@ impl<T: ValueType> ValueType for NullableType<T> {`
`69`	`69`	`has_null: *has_null,`
`70`	`70`	`value: Some(Box::new(T::try_downcast_domain(value)?)),`
`71`	`71`	`}),`
	`72`	`+ Domain::Nullable(NullableDomain {`
	`73`	`+ has_null,`
	`74`	`+ value: None,`
	`75`	`+ }) => Some(NullableDomain {`
	`76`	`+ has_null: *has_null,`
	`77`	`+ value: None,`
	`78`	`+ }),`
`72`	`79`	`_ => None,`
`73`	`80`	`}`
`74`	`81`	`}`