diff --git a/dbms/src/Flash/Coprocessor/DAGUtils.cpp b/dbms/src/Flash/Coprocessor/DAGUtils.cpp index c83f029c7f9..4b1d3546406 100644 --- a/dbms/src/Flash/Coprocessor/DAGUtils.cpp +++ b/dbms/src/Flash/Coprocessor/DAGUtils.cpp @@ -610,7 +610,7 @@ const std::unordered_map scalar_func_map({ {tipb::ScalarFuncSig::Concat, "tidbConcat"}, {tipb::ScalarFuncSig::ConcatWS, "tidbConcatWS"}, //{tipb::ScalarFuncSig::Convert, "cast"}, - //{tipb::ScalarFuncSig::Elt, "cast"}, + {tipb::ScalarFuncSig::Elt, "elt"}, //{tipb::ScalarFuncSig::ExportSet3Arg, "cast"}, //{tipb::ScalarFuncSig::ExportSet4Arg, "cast"}, //{tipb::ScalarFuncSig::ExportSet5Arg, "cast"}, diff --git a/dbms/src/Functions/FunctionHelpers.h b/dbms/src/Functions/FunctionHelpers.h index 1c4759f531d..e4fa79d9a8a 100644 --- a/dbms/src/Functions/FunctionHelpers.h +++ b/dbms/src/Functions/FunctionHelpers.h @@ -55,6 +55,16 @@ bool checkColumn(const IColumn * column) return checkAndGetColumn(column); } +template +const Type * checkAndGetNestedColumn(const IColumn * column) +{ + if (!column || !column->isColumnNullable()) + return {}; + + const auto * data_column = &static_cast(column)->getNestedColumn(); + + return checkAndGetColumn(data_column); +} template const ColumnConst * checkAndGetColumnConst(const IColumn * column, bool maybe_nullable_column = false) @@ -62,7 +72,7 @@ const ColumnConst * checkAndGetColumnConst(const IColumn * column, bool maybe_nu if (!column || !column->isColumnConst()) return {}; - const ColumnConst * res = static_cast(column); + const auto * res = static_cast(column); const auto * data_column = &res->getDataColumn(); if (maybe_nullable_column && data_column->isColumnNullable()) diff --git a/dbms/src/Functions/FunctionsString.cpp b/dbms/src/Functions/FunctionsString.cpp index aac766e3e53..b9b4c570ef5 100644 --- a/dbms/src/Functions/FunctionsString.cpp +++ b/dbms/src/Functions/FunctionsString.cpp @@ -29,15 +29,11 @@ #include #include #include -#include #include #include -#include -#include #include #include -#include namespace DB { @@ -5415,8 +5411,208 @@ class FunctionBin : public IFunction throw Exception(fmt::format("Illegal argument of function {}", getName()), ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT); } } +}; + +class FunctionElt : public IFunction +{ +public: + static constexpr auto name = "elt"; + + static FunctionPtr create(const Context & /*context*/) + { + return std::make_shared(); + } + + String getName() const override { return name; } + size_t getNumberOfArguments() const override { return 0; } + bool isVariadic() const override { return true; } + + bool useDefaultImplementationForNulls() const override { return false; } + bool useDefaultImplementationForConstants() const override { return true; } + + DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override + { + if (arguments.size() < 2) + throw Exception( + fmt::format("Number of arguments for function {} doesn't match: passed {}, should be at least 2.", getName(), arguments.size()), + ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH); + + auto first_argument = removeNullable(arguments[0]); + if (!first_argument->isInteger()) + throw Exception( + fmt::format("Illegal type {} of first argument of function {}", first_argument->getName(), getName()), + ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT); + + for (const auto arg_idx : ext::range(1, arguments.size())) + { + const auto arg = removeNullable(arguments[arg_idx]); + if (!arg->isString()) + throw Exception( + fmt::format("Illegal type {} of argument {} of function {}", arg->getName(), arg_idx + 1, getName()), + ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT); + } + + return makeNullable(std::make_shared()); + } + + void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result) const override + { + if (executeElt(block, arguments, result) + || executeElt(block, arguments, result) + || executeElt(block, arguments, result) + || executeElt(block, arguments, result) + || executeElt(block, arguments, result) + || executeElt(block, arguments, result) + || executeElt(block, arguments, result) + || executeElt(block, arguments, result)) + { + return; + } + else + { + throw Exception(fmt::format("Illegal argument of function {}", getName()), ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT); + } + } private: + using NullMapMutablePtr = COWPtrHelper>::MutablePtr; + + template + static bool executeElt(Block & block, const ColumnNumbers & arguments, size_t result) + { + const auto * col_arg0 = block.getByPosition(arguments[0]).column.get(); + + if (const auto * col = checkAndGetColumnConst>(col_arg0, true)) + { + return constColumn(col, block, arguments, result); + } + else + { + return vectorColumn(col_arg0, block, arguments, result); + } + } + + static void fillResultColumnNull(ColumnPtr & dst, size_t nrow) + { + dst = DataTypeNullable(std::make_shared()).createColumnConst(nrow, {}); + } + + static void fillResultColumnFromOther(ColumnPtr & dst, const ColumnPtr & src) + { + dst = makeNullable(src->cloneResized(src->size())); + } + + /// fill the ith element of result column from the ith element of another column + /// Note that for efficiency purpose, following preconditions should be satisfied + /// 1. res_null_map should already have enough size to contain the ith element and its default value should be 0 + /// 2. res_offsets should already be resized to be able to contain the ith element, therefore no `push_back` can be used + /// 3. res_chars should **not** be sized already for ith element, but can be reserved to have enough space + static void fillResultColumnEntry(NullMapMutablePtr & res_null_map, ColumnString::Chars_t & res_chars, IColumn::Offsets & res_offsets, const ColumnPtr & src, const size_t dsti) + { + if (src->isNullAt(dsti)) + { + res_null_map->getData()[dsti] = true; + res_chars.push_back(0); + res_offsets[dsti] = dsti == 0 ? 1 : (res_offsets[dsti - 1] + 1); + return; + } + + /// no need to set res_null_map, since its default value is 0 + + /// src col might be ColumnConst(Nullable(ColumnString)) or ColumnCost(ColumnString) or Nullable(ColumnString) or ColumnString + /// if it is ColumnConst(...) then we should treat its first element as the ith element + size_t srci = dsti; + const auto * col_nullable_str = src->isColumnConst() + ? (srci = 0, checkAndGetColumnConst(src.get(), true)->getDataColumnPtr().get()) + : src.get(); + + const auto * col_str = col_nullable_str->isColumnNullable() + ? checkAndGetNestedColumn(col_nullable_str) + : checkAndGetColumn(col_nullable_str); + + const auto & src_data = col_str->getChars(); + const auto & src_offsets = col_str->getOffsets(); + + const auto start_offset = StringUtil::offsetAt(src_offsets, srci); + const auto str_size = StringUtil::sizeAt(src_offsets, srci); + + const size_t old_size = res_chars.size(); + const size_t new_size = old_size + str_size; + + res_chars.resize(new_size); + memcpy(&res_chars[old_size], &src_data[start_offset], str_size); + res_offsets[dsti] = new_size; + } + + template + static bool constColumn(const ColumnConst * col, Block & block, const ColumnNumbers & arguments, size_t result) + { + const auto nrow = col->size(); + + if (col->onlyNull()) + { + fillResultColumnNull(block.getByPosition(result).column, nrow); + return true; + } + + /// get the first argument from the const column which still might be nullable + const auto arg0 = col->getDataColumnPtr()->isColumnNullable() + ? checkAndGetNestedColumn>(col->getDataColumnPtr().get())->getInt(0) + : col->getInt(0); + + if (arg0 < 1 || arg0 >= static_cast(arguments.size())) + { + fillResultColumnNull(block.getByPosition(result).column, nrow); + } + else + { + fillResultColumnFromOther(block.getByPosition(result).column, block.getByPosition(arg0).column); + } + return true; + } + + template + static bool vectorColumn(const IColumn * col, Block & block, const ColumnNumbers & arguments, size_t result) + { + const auto narg = arguments.size(); + const auto nrow = col->size(); + const auto col_arg0 = col->isColumnNullable() + ? checkAndGetNestedColumn>(col) + : checkAndGetColumn>(col); + + if (!col_arg0) + { + return false; + } + + const auto & arg0_vec = col_arg0->getData(); + + auto res_null_map = ColumnUInt8::create(nrow, false); + auto res_col = ColumnString::create(); + auto & res_chars = res_col->getChars(); + auto & res_offsets = res_col->getOffsets(); + + res_offsets.resize_fill(nrow); + + for (size_t i = 0; i < nrow; ++i) + { + const auto arg0 = arg0_vec[i]; + + if (col_arg0->isNullAt(i) || arg0 < 1 || static_cast(arg0) >= static_cast(narg)) + { + res_null_map->getData()[i] = true; + res_chars.push_back(0); + res_offsets[i] = i == 0 ? 1 : (res_offsets[i - 1] + 1); + } + else + { + fillResultColumnEntry(res_null_map, res_chars, res_offsets, block.getByPosition(arguments[arg0]).column, i); + } + } + + block.getByPosition(result).column = ColumnNullable::create(std::move(res_col), std::move(res_null_map)); + return true; + } }; // clang-format off @@ -5507,5 +5703,6 @@ void registerFunctionsString(FunctionFactory & factory) factory.registerFunction(); factory.registerFunction(); factory.registerFunction(); + factory.registerFunction(); } } // namespace DB diff --git a/dbms/src/Functions/tests/gtest_strings_elt.cpp b/dbms/src/Functions/tests/gtest_strings_elt.cpp new file mode 100644 index 00000000000..4dfde76f76a --- /dev/null +++ b/dbms/src/Functions/tests/gtest_strings_elt.cpp @@ -0,0 +1,265 @@ +// Copyright 2022 PingCAP, Ltd. +// +// Licensed 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 +#include + +#include + +namespace DB +{ +namespace tests +{ +class TestFunctionElt : public DB::tests::FunctionTest +{ +}; + +#define ASSERT_ELT(expected, t1, t2, ...) \ + ASSERT_COLUMN_EQ(expected, executeFunction("elt", t1, t2, __VA_ARGS__)) + +TEST_F(TestFunctionElt, BoundaryIdx) +try +{ + constexpr size_t ndummy_col = 6; + const ColumnWithTypeAndName dummy_col0 = createColumn>({"abc", "123", "", {}, "", ""}); + const ColumnWithTypeAndName dummy_col1 = createConstColumn>(ndummy_col, "def"); + + constexpr size_t nboundary_test = 6; + const int boundary_values[nboundary_test] = { + 0, + -1, + -2, + -100, + 3, + 100}; + + /// elt returns null if the first argument is less than 1, greater than the number of string arguments, or NULL + + /// const idx + for (const auto i : ext::range(0, nboundary_test)) + { + ASSERT_ELT(createConstColumn>(ndummy_col, {}), + createConstColumn>(ndummy_col, boundary_values[i]), + dummy_col0, + dummy_col1); + + ASSERT_ELT(createConstColumn>(ndummy_col, {}), + createConstColumn(ndummy_col, boundary_values[i]), + dummy_col0, + dummy_col1); + } + + /// vector idx + ASSERT_ELT(createColumn>({{}, {}, {}, {}, {}, {}}), + createColumn>({0, -1, -2, -100, 3, 100}), + dummy_col0, + dummy_col1); +} +CATCH + +TEST_F(TestFunctionElt, NullArgument) +try +{ + constexpr size_t nrow = 5; + + /// const null idx + ASSERT_ELT(createConstColumn>(nrow, {}), + createConstColumn>(nrow, {}), + createColumn>({"abc", "123", "", {}, ""}), + createConstColumn>(nrow, "def")); + + /// vector null idx + ASSERT_ELT(createColumn>({{}, {}, {}, {}, {}}), + createColumn>({{}, {}, {}, {}, {}}), + createColumn>({"abc", "123", "", {}, ""}), + createConstColumn>(nrow, "def")); + + /// const non-null idx x const null arg + ASSERT_ELT(createConstColumn>(nrow, {}), + createConstColumn>(nrow, 1), + createConstColumn>(nrow, {}), + createColumn>({"def", "321", {}, "", ""})); + + /// vector non-null idx x const null arg + ASSERT_ELT(createColumn>({{}, {}, {}, {}, {}}), + createColumn>({1, 1, 1, 1, 1}), + createConstColumn>(nrow, {}), + createColumn>({"def", "321", {}, "", ""})); + + /// const non-null idx x null vector arg + ASSERT_ELT(createColumn>({{}, {}, {}, {}, {}}), + createConstColumn>(nrow, 1), + createColumn>({{}, {}, {}, {}, {}}), + createColumn>({"def", "321", {}, "", ""})); + + /// vector non-null idx x null vector arg + ASSERT_ELT(createColumn>({{}, {}, {}, {}, {}}), + createColumn>({1, 1, 1, 1, 1}), + createColumn>({{}, {}, {}, {}, {}}), + createColumn>({"def", "321", {}, "", ""})); + + /// const non-null idx x nullable arg + ASSERT_ELT(createColumn>({"abc", "123", "", {}, ""}), + createConstColumn>(nrow, 1), + createColumn>({"abc", "123", "", {}, ""}), + createColumn>({"def", "321", {}, "", ""})); + + /// vector idx x nullable arg + ASSERT_ELT(createColumn>({"abc", "321", {}, {}, {}}), + createColumn>({1, 2, 2, 1, {}}), + createColumn>({"abc", "123", "", {}, ""}), + createColumn>({"def", "321", {}, "", ""})); +} +CATCH + +TEST_F(TestFunctionElt, AllTypeIdx) +try +{ + constexpr size_t nrow = 8; + + /// const idx x vector + ASSERT_ELT(createColumn>({"abc", "123", "", "", {}, "", "a", "c"}), + createConstColumn>(nrow, 1), + createColumn>({"abc", "123", "", "", {}, "", "a", "c"}), + createColumn>({"def", "321", {}, "", "", "", "b", "d"})); + + ASSERT_ELT(createColumn>({"def", "321", {}, "", "", "", "b", "d"}), + createConstColumn>(nrow, 2), + createColumn>({"abc", "123", "", "", {}, "", "a", "c"}), + createColumn>({"def", "321", {}, "", "", "", "b", "d"})); + + ASSERT_ELT(createColumn>({"abc", "123", "", "", {}, "", "a", "c"}), + createConstColumn>(nrow, 1), + createColumn>({"abc", "123", "", "", {}, "", "a", "c"}), + createColumn>({"def", "321", {}, "", "", "", "b", "d"})); + + ASSERT_ELT(createColumn>({"def", "321", {}, "", "", "", "b", "d"}), + createConstColumn>(nrow, 2), + createColumn>({"abc", "123", "", "", {}, "", "a", "c"}), + createColumn>({"def", "321", {}, "", "", "", "b", "d"})); + + ASSERT_ELT(createColumn>({"abc", "123", "", "", {}, "", "a", "c"}), + createConstColumn>(nrow, 1), + createColumn>({"abc", "123", "", "", {}, "", "a", "c"}), + createColumn>({"def", "321", {}, "", "", "", "b", "d"})); + + ASSERT_ELT(createColumn>({"def", "321", {}, "", "", "", "b", "d"}), + createConstColumn>(nrow, 2), + createColumn>({"abc", "123", "", "", {}, "", "a", "c"}), + createColumn>({"def", "321", {}, "", "", "", "b", "d"})); + + ASSERT_ELT(createColumn>({"abc", "123", "", "", {}, "", "a", "c"}), + createConstColumn>(nrow, 1), + createColumn>({"abc", "123", "", "", {}, "", "a", "c"}), + createColumn>({"def", "321", {}, "", "", "", "b", "d"})); + + ASSERT_ELT(createColumn>({"def", "321", {}, "", "", "", "b", "d"}), + createConstColumn>(nrow, 2), + createColumn>({"abc", "123", "", "", {}, "", "a", "c"}), + createColumn>({"def", "321", {}, "", "", "", "b", "d"})); + + /// const idx x const + ASSERT_ELT(createConstColumn>(nrow, "abc"), + createConstColumn>(nrow, 1), + createConstColumn>(nrow, "abc"), + createConstColumn>(nrow, "def")); + + ASSERT_ELT(createConstColumn>(nrow, "def"), + createConstColumn>(nrow, 2), + createConstColumn>(nrow, "abc"), + createConstColumn>(nrow, "def")); + + ASSERT_ELT(createConstColumn>(nrow, "abc"), + createConstColumn>(nrow, 1), + createConstColumn>(nrow, "abc"), + createConstColumn>(nrow, "def")); + + ASSERT_ELT(createConstColumn>(nrow, "def"), + createConstColumn>(nrow, 2), + createConstColumn>(nrow, "abc"), + createConstColumn>(nrow, "def")); + + ASSERT_ELT(createConstColumn>(nrow, "abc"), + createConstColumn>(nrow, 1), + createConstColumn>(nrow, "abc"), + createConstColumn>(nrow, "def")); + + ASSERT_ELT(createConstColumn>(nrow, "def"), + createConstColumn>(nrow, 2), + createConstColumn>(nrow, "abc"), + createConstColumn>(nrow, "def")); + + ASSERT_ELT(createConstColumn>(nrow, "abc"), + createConstColumn>(nrow, 1), + createConstColumn>(nrow, "abc"), + createConstColumn>(nrow, "def")); + + ASSERT_ELT(createConstColumn>(nrow, "def"), + createConstColumn>(nrow, 2), + createConstColumn>(nrow, "abc"), + createConstColumn>(nrow, "def")); + + /// vector idx x vector/const + ASSERT_ELT(createColumn>({"abc", "321", "", {}, "xxx", "", {}, {}}), + createColumn>({1, 2, 1, 100, 3, 1, 0, -2}), + createColumn>({"abc", "123", "", "", {}, "", "a", "c"}), + createColumn>({"def", "321", {}, "", "", "", "b", "d"}), + createConstColumn>(nrow, "xxx")); + + ASSERT_ELT(createColumn>({"abc", "321", "", {}, "xxx", "", {}, "c"}), + createColumn>({1, 2, 1, 100, 3, 1, 0, 1}), + createColumn>({"abc", "123", "", "", {}, "", "a", "c"}), + createColumn>({"def", "321", {}, "", "", "", "b", "d"}), + createConstColumn>(nrow, "xxx")); + + ASSERT_ELT(createColumn>({"abc", "321", "xxx", {}, "xxx", "", "a", {}}), + createColumn>({1, 2, 3, 100, 3, 2, 1, -1}), + createColumn>({"abc", "123", "", "", {}, "", "a", "c"}), + createColumn>({"def", "321", {}, "", "", "", "b", "d"}), + createConstColumn>(nrow, "xxx")); + + ASSERT_ELT(createColumn>({"abc", "321", "", {}, "xxx", "", {}, "c"}), + createColumn>({1, 2, 1, 100, 3, 1, 0, 1}), + createColumn>({"abc", "123", "", "", {}, "", "a", "c"}), + createColumn>({"def", "321", {}, "", "", "", "b", "d"}), + createConstColumn>(nrow, "xxx")); + + ASSERT_ELT(createColumn>({"abc", "321", "xxx", {}, "xxx", "", "a", {}}), + createColumn>({1, 2, 3, 100, 3, 2, 1, -1}), + createColumn>({"abc", "123", "", "", {}, "", "a", "c"}), + createColumn>({"def", "321", {}, "", "", "", "b", "d"}), + createConstColumn>(nrow, "xxx")); + + ASSERT_ELT(createColumn>({"abc", "321", "", {}, "xxx", "", {}, "c"}), + createColumn>({1, 2, 1, 100, 3, 1, 0, 1}), + createColumn>({"abc", "123", "", "", {}, "", "a", "c"}), + createColumn>({"def", "321", {}, "", "", "", "b", "d"}), + createConstColumn>(nrow, "xxx")); + + ASSERT_ELT(createColumn>({"abc", "321", "xxx", {}, "xxx", "", "a", {}}), + createColumn>({1, 2, 3, 100, 3, 2, 1, -1}), + createColumn>({"abc", "123", "", "", {}, "", "a", "c"}), + createColumn>({"def", "321", {}, "", "", "", "b", "d"}), + createConstColumn>(nrow, "xxx")); + + ASSERT_ELT(createColumn>({"abc", "321", "", {}, "xxx", "", {}, "c"}), + createColumn>({1, 2, 1, 100, 3, 1, 0, 1}), + createColumn>({"abc", "123", "", "", {}, "", "a", "c"}), + createColumn>({"def", "321", {}, "", "", "", "b", "d"}), + createConstColumn>(nrow, "xxx")); +} +CATCH + +} // namespace tests +} // namespace DB diff --git a/tests/fullstack-test/expr/elt.test b/tests/fullstack-test/expr/elt.test new file mode 100644 index 00000000000..46c675befed --- /dev/null +++ b/tests/fullstack-test/expr/elt.test @@ -0,0 +1,82 @@ +# Copyright 2022 PingCAP, Ltd. +# +# Licensed 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. + +mysql> drop table if exists test.t; +mysql> create table test.t (a int, b varchar(200), c varchar(200)); +mysql> insert into test.t values(-1, 'abc', 'def'); +mysql> insert into test.t values(0, '▲α▼βγ➨δε☎ζη✂θι€κλ♫μν✓ξο✚πρ℉στ♥υφ♖χψ♘ω★σ✕', 'hello world'); +mysql> insert into test.t values(1, 'թփձջրչճժծքոեռտըւիօպասդֆգհյկլխզղցվբնմշ', 'world hello'); +mysql> insert into test.t values(2, NULL, '11111'); +mysql> insert into test.t values(3, NULL, '11111'); +mysql> alter table test.t set tiflash replica 1; + +func> wait_table test t + +mysql> set tidb_enforce_mpp=1; set @@session.tidb_isolation_read_engines = "tiflash"; select elt(a, b, c) as v from test.t; +v +NULL +NULL +թփձջրչճժծքոեռտըւիօպասդֆգհյկլխզղցվբնմշ +11111 +NULL + +mysql> set tidb_enforce_mpp=1; set @@session.tidb_isolation_read_engines = "tiflash"; select elt(NULL, b, c) as v from test.t; +v +NULL +NULL +NULL +NULL +NULL + +mysql> set tidb_enforce_mpp=1; set @@session.tidb_isolation_read_engines = "tiflash"; select elt(-1, b, c) as v from test.t; +v +NULL +NULL +NULL +NULL +NULL + +mysql> set tidb_enforce_mpp=1; set @@session.tidb_isolation_read_engines = "tiflash"; select elt(0, b, c) as v from test.t; +v +NULL +NULL +NULL +NULL +NULL + +mysql> set tidb_enforce_mpp=1; set @@session.tidb_isolation_read_engines = "tiflash"; select elt(1, b, c) as v from test.t; +v +abc +▲α▼βγ➨δε☎ζη✂θι€κλ♫μν✓ξο✚πρ℉στ♥υφ♖χψ♘ω★σ✕ +թփձջրչճժծքոեռտըւիօպասդֆգհյկլխզղցվբնմշ +NULL +NULL + +mysql> set tidb_enforce_mpp=1; set @@session.tidb_isolation_read_engines = "tiflash"; select elt(2, b, c) as v from test.t; +v +def +hello world +world hello +11111 +11111 + +mysql> set tidb_enforce_mpp=1; set @@session.tidb_isolation_read_engines = "tiflash"; select elt(3, b, c) as v from test.t; +v +NULL +NULL +NULL +NULL +NULL + +mysql> drop table if exists test.t