This is an automated email from the ASF dual-hosted git repository. starocean999 pushed a commit to branch branch-3.0 in repository https://gitbox.apache.org/repos/asf/doris.git
The following commit(s) were added to refs/heads/branch-3.0 by this push: new 9ca5ff4c108 [branch-3.0](function) Support function assert_true (#39393) (#41898) 9ca5ff4c108 is described below commit 9ca5ff4c1089e0831b96a39eddc047a320cc2e9d Author: zclllhhjj <zhaochan...@selectdb.com> AuthorDate: Wed Oct 16 10:15:55 2024 +0800 [branch-3.0](function) Support function assert_true (#39393) (#41898) pick https://github.com/apache/doris/pull/39393 --- be/src/common/exception.h | 2 - be/src/vec/data_types/data_type_number.h | 1 + be/src/vec/functions/function_assert_true.cpp | 97 ++++++++++++++++++++++ be/src/vec/functions/simple_function_factory.h | 2 + .../doris/catalog/BuiltinScalarFunctions.java | 2 + .../expressions/functions/scalar/AssertTrue.java | 76 +++++++++++++++++ .../expressions/visitor/ScalarFunctionVisitor.java | 5 ++ .../conditional_functions/test_assert_true.out | 20 +++++ .../conditional_functions/test_assert_true.groovy | 79 ++++++++++++++++++ 9 files changed, 282 insertions(+), 2 deletions(-) diff --git a/be/src/common/exception.h b/be/src/common/exception.h index 37d6afe937e..8d35ae4cb5f 100644 --- a/be/src/common/exception.h +++ b/be/src/common/exception.h @@ -20,10 +20,8 @@ #include <fmt/format.h> #include <gen_cpp/Status_types.h> -#include <cstdint> #include <exception> #include <memory> -#include <ostream> #include <string> #include <string_view> #include <utility> diff --git a/be/src/vec/data_types/data_type_number.h b/be/src/vec/data_types/data_type_number.h index aa3a9ca21e5..e1612c029dc 100644 --- a/be/src/vec/data_types/data_type_number.h +++ b/be/src/vec/data_types/data_type_number.h @@ -50,6 +50,7 @@ using DataTypeInt64 = DataTypeNumber<Int64>; using DataTypeInt128 = DataTypeNumber<Int128>; using DataTypeFloat32 = DataTypeNumber<Float32>; using DataTypeFloat64 = DataTypeNumber<Float64>; +using DataTypeBool = DataTypeUInt8; template <typename DataType> constexpr bool IsDataTypeNumber = false; diff --git a/be/src/vec/functions/function_assert_true.cpp b/be/src/vec/functions/function_assert_true.cpp new file mode 100644 index 00000000000..71deeab32a5 --- /dev/null +++ b/be/src/vec/functions/function_assert_true.cpp @@ -0,0 +1,97 @@ +// 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. + +#include <memory> + +#include "common/exception.h" +#include "common/status.h" +#include "vec/columns/column.h" +#include "vec/columns/column_const.h" +#include "vec/columns/column_nullable.h" +#include "vec/columns/columns_number.h" +#include "vec/core/block.h" +#include "vec/core/column_numbers.h" +#include "vec/core/column_with_type_and_name.h" +#include "vec/core/types.h" +#include "vec/data_types/data_type.h" +#include "vec/data_types/data_type_number.h" +#include "vec/functions/function.h" +#include "vec/functions/function_helpers.h" +#include "vec/functions/simple_function_factory.h" + +namespace doris { +class FunctionContext; +} // namespace doris + +namespace doris::vectorized { + +class FunctionAssertTrue : public IFunction { +public: + static constexpr auto name = "assert_true"; + + static FunctionPtr create() { return std::make_shared<FunctionAssertTrue>(); } + + String get_name() const override { return name; } + + size_t get_number_of_arguments() const override { return 2; } + + Status open(FunctionContext* context, FunctionContext::FunctionStateScope scope) override { + if (!context->is_col_constant(1)) [[unlikely]] { + return Status::InvalidArgument("assert_true only accept constant for 2nd argument"); + } + return Status::OK(); + } + + DataTypePtr get_return_type_impl(const DataTypes& arguments) const override { + return std::make_shared<DataTypeBool>(); + } + + bool use_default_implementation_for_nulls() const final { return false; } + ColumnNumbers get_arguments_that_are_always_constant() const final { return {1}; } + + // column2 is const, so in default logic column1 is no way const. + Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments, + size_t result, size_t input_rows_count) const override { + std::string errmsg = + assert_cast<const ColumnConst&>(*block.get_by_position(arguments[1]).column) + .get_data_at(0) + .to_string(); + + ColumnPtr col_ptr = block.get_by_position(arguments[0]).column; + if (const auto* col_nullable = check_and_get_column<ColumnNullable>(col_ptr.get())) { + if (col_nullable->has_null()) { + throw doris::Exception(Status::InvalidArgument(errmsg)); + } + col_ptr = col_nullable->get_nested_column_ptr(); + } + const auto& data = assert_cast<const ColumnBool*>(col_ptr.get())->get_data(); + + for (int i = 0; i < input_rows_count; i++) { + if (!data[i]) [[unlikely]] { + throw doris::Exception(Status::InvalidArgument(errmsg)); + } + } + + block.replace_by_position(result, ColumnBool::create(input_rows_count, true)); + return Status::OK(); + } +}; + +void register_function_assert_true(SimpleFunctionFactory& factory) { + factory.register_function<FunctionAssertTrue>(); +} +} // namespace doris::vectorized diff --git a/be/src/vec/functions/simple_function_factory.h b/be/src/vec/functions/simple_function_factory.h index c07b885c70c..407855484ff 100644 --- a/be/src/vec/functions/simple_function_factory.h +++ b/be/src/vec/functions/simple_function_factory.h @@ -105,6 +105,7 @@ void register_function_tokenize(SimpleFunctionFactory& factory); void register_function_url(SimpleFunctionFactory& factory); void register_function_ip(SimpleFunctionFactory& factory); void register_function_multi_match(SimpleFunctionFactory& factory); +void register_function_assert_true(SimpleFunctionFactory& factory); class SimpleFunctionFactory { using Creator = std::function<FunctionBuilderPtr()>; @@ -284,6 +285,7 @@ public: register_function_ignore(instance); register_function_variant_element(instance); register_function_multi_match(instance); + register_function_assert_true(instance); }); return instance; } diff --git a/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinScalarFunctions.java b/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinScalarFunctions.java index 1b49917b85f..587805ba3a7 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinScalarFunctions.java +++ b/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinScalarFunctions.java @@ -71,6 +71,7 @@ import org.apache.doris.nereids.trees.expressions.functions.scalar.ArrayZip; import org.apache.doris.nereids.trees.expressions.functions.scalar.ArraysOverlap; import org.apache.doris.nereids.trees.expressions.functions.scalar.Ascii; import org.apache.doris.nereids.trees.expressions.functions.scalar.Asin; +import org.apache.doris.nereids.trees.expressions.functions.scalar.AssertTrue; import org.apache.doris.nereids.trees.expressions.functions.scalar.Atan; import org.apache.doris.nereids.trees.expressions.functions.scalar.Atan2; import org.apache.doris.nereids.trees.expressions.functions.scalar.AutoPartitionName; @@ -528,6 +529,7 @@ public class BuiltinScalarFunctions implements FunctionHelper { scalar(ArraysOverlap.class, "arrays_overlap"), scalar(Ascii.class, "ascii"), scalar(Asin.class, "asin"), + scalar(AssertTrue.class, "assert_true"), scalar(Atan.class, "atan"), scalar(Atan2.class, "atan2"), scalar(AutoPartitionName.class, "auto_partition_name"), diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/AssertTrue.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/AssertTrue.java new file mode 100644 index 00000000000..a4bffd9903a --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/AssertTrue.java @@ -0,0 +1,76 @@ +// 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. + +package org.apache.doris.nereids.trees.expressions.functions.scalar; + +import org.apache.doris.catalog.FunctionSignature; +import org.apache.doris.nereids.exceptions.AnalysisException; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.functions.AlwaysNotNullable; +import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature; +import org.apache.doris.nereids.trees.expressions.shape.BinaryExpression; +import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor; +import org.apache.doris.nereids.types.BooleanType; +import org.apache.doris.nereids.types.VarcharType; + +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; + +import java.util.List; + +/** + * ScalarFunction 'assert_true'. + */ +public class AssertTrue extends ScalarFunction + implements BinaryExpression, ExplicitlyCastableSignature, AlwaysNotNullable { + + public static final List<FunctionSignature> SIGNATURES = ImmutableList.of( + FunctionSignature.ret(BooleanType.INSTANCE).args(BooleanType.INSTANCE, VarcharType.SYSTEM_DEFAULT)); + + /** + * constructor with 2 argument. + */ + public AssertTrue(Expression arg0, Expression arg1) { + super("assert_true", arg0, arg1); + } + + /** + * withChildren. + */ + @Override + public AssertTrue withChildren(List<Expression> children) { + Preconditions.checkArgument(children.size() == 2); + return new AssertTrue(children.get(0), children.get(1)); + } + + @Override + public List<FunctionSignature> getSignatures() { + return SIGNATURES; + } + + @Override + public void checkLegalityAfterRewrite() { + if (!getArgument(1).isLiteral()) { + throw new AnalysisException("assert_true only accept constant for 2nd argument"); + } + } + + @Override + public <R, C> R accept(ExpressionVisitor<R, C> visitor, C context) { + return visitor.visitAssertTrue(this, context); + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java index a905ae28d8f..56f1e724650 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java @@ -78,6 +78,7 @@ import org.apache.doris.nereids.trees.expressions.functions.scalar.ArrayZip; import org.apache.doris.nereids.trees.expressions.functions.scalar.ArraysOverlap; import org.apache.doris.nereids.trees.expressions.functions.scalar.Ascii; import org.apache.doris.nereids.trees.expressions.functions.scalar.Asin; +import org.apache.doris.nereids.trees.expressions.functions.scalar.AssertTrue; import org.apache.doris.nereids.trees.expressions.functions.scalar.Atan; import org.apache.doris.nereids.trees.expressions.functions.scalar.Atan2; import org.apache.doris.nereids.trees.expressions.functions.scalar.AutoPartitionName; @@ -696,6 +697,10 @@ public interface ScalarFunctionVisitor<R, C> { return visitScalarFunction(asin, context); } + default R visitAssertTrue(AssertTrue assertTrue, C context) { + return visitScalarFunction(assertTrue, context); + } + default R visitAtan(Atan atan, C context) { return visitScalarFunction(atan, context); } diff --git a/regression-test/data/query_p0/sql_functions/conditional_functions/test_assert_true.out b/regression-test/data/query_p0/sql_functions/conditional_functions/test_assert_true.out new file mode 100644 index 00000000000..6066a2fb852 --- /dev/null +++ b/regression-test/data/query_p0/sql_functions/conditional_functions/test_assert_true.out @@ -0,0 +1,20 @@ +-- This file is automatically generated. You should know what you did if you want to edit this +-- !sql1 -- +true + +-- !sql2 -- +true + +-- !sql3 -- +true +true +true + +-- !sql4 -- +true +true +true + +-- !sql -- +true + diff --git a/regression-test/suites/query_p0/sql_functions/conditional_functions/test_assert_true.groovy b/regression-test/suites/query_p0/sql_functions/conditional_functions/test_assert_true.groovy new file mode 100644 index 00000000000..f84044ee507 --- /dev/null +++ b/regression-test/suites/query_p0/sql_functions/conditional_functions/test_assert_true.groovy @@ -0,0 +1,79 @@ +// 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. + +suite("test_assert_true") { + sql " drop table if exists assert_true_not_null " + sql """ + create table assert_true_not_null( + k0 boolean not null, + k1 int + ) + DISTRIBUTED BY HASH(`k0`) BUCKETS auto + properties("replication_num" = "1"); + """ + sql " insert into assert_true_not_null values (true, 1);" + qt_sql1 """ select assert_true(k0, cast(123 as varchar)) from assert_true_not_null; """ + qt_sql2 """ select assert_true(k0, "nn123") from assert_true_not_null; """ + + sql " insert into assert_true_not_null values (false, 2);" + test { + sql """ select assert_true(k0, cast(123 as varchar)) from assert_true_not_null; """ + exception "123" + } + test { + sql """ select assert_true(k0, "nn123") from assert_true_not_null; """ + exception "nn123" + } + test { + sql """ select assert_true(1, k1) from assert_true_not_null; """ + exception "assert_true only accept constant for 2nd argument" + } + + + + sql " drop table if exists assert_true_null " + sql """ + create table assert_true_null( + k0 boolean null + ) + DISTRIBUTED BY HASH(`k0`) BUCKETS auto + properties("replication_num" = "1"); + """ + sql " insert into assert_true_null values (true), (true), (true);" + qt_sql3 """ select assert_true(k0, cast(123 as varchar)) from assert_true_null; """ + qt_sql4 """ select assert_true(k0, "nn123") from assert_true_null; """ + + sql " insert into assert_true_null values (null);" + test { + sql """ select assert_true(k0, cast(123 as varchar)) from assert_true_null; """ + exception "123" + } + test { + sql """ select assert_true(k0, "nn123") from assert_true_null; """ + exception "nn123" + } + + qt_sql """ select assert_true(1, "wrong"); """ + test { + sql """ select assert_true(0, nullable("wrong")); """ + exception "assert_true only accept constant for 2nd argument" + } + test { + sql """ select assert_true(0, "wrong"); """ + exception "wrong" + } +} --------------------------------------------------------------------- To unsubscribe, e-mail: commits-unsubscr...@doris.apache.org For additional commands, e-mail: commits-h...@doris.apache.org