fix error when first argument of replace function is const (#9615)

dbms/src/Functions/FunctionsStringReplace.h

@@ -0,0 +1,365 @@
+// Copyright 2023 PingCAP, Inc.
+//
+// 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.
+
+#pragma once
+
+#include <Columns/ColumnConst.h>
+#include <Columns/ColumnFixedString.h>
+#include <Columns/ColumnString.h>
+#include <Columns/IColumn.h>
+#include <Common/Exception.h>
+#include <DataTypes/DataTypeString.h>
+#include <DataTypes/DataTypesNumber.h>
+#include <Functions/FunctionHelpers.h>
+#include <Functions/GatherUtils/Algorithms.h>
+#include <Functions/GatherUtils/Sources.h>
+#include <Functions/IFunction.h>
+
+namespace DB
+{
+
+namespace ErrorCodes
+{
+extern const int ILLEGAL_COLUMN;
+} // namespace ErrorCodes
+
+template <typename Impl, typename Name>
+class FunctionStringReplace : public IFunction
+{
+public:
+    static constexpr auto name = Name::name;
+    static FunctionPtr create(const Context &) { return std::make_shared<FunctionStringReplace>(); }
+
+    String getName() const override { return name; }
+
+    size_t getNumberOfArguments() const override { return 3; }
+
+    bool isVariadic() const override { return false; }
+    bool useDefaultImplementationForConstants() const override { return true; }
+
+    DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
+    {
+        if (!arguments[0]->isStringOrFixedString())
+            throw Exception(
+                "Illegal type " + arguments[0]->getName() + " of first argument of function " + getName(),
+                ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
+
+        if (!arguments[1]->isStringOrFixedString())
+            throw Exception(
+                "Illegal type " + arguments[1]->getName() + " of second argument of function " + getName(),
+                ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
+
+        if (!arguments[2]->isStringOrFixedString())
+            throw Exception(
+                "Illegal type " + arguments[2]->getName() + " of third argument of function " + getName(),
+                ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
+
+        return std::make_shared<DataTypeString>();
+    }
+
+    void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result) const override
+    {
+        ColumnPtr column_src = block.getByPosition(arguments[0]).column;
+        ColumnPtr column_needle = block.getByPosition(arguments[1]).column;
+        ColumnPtr column_replacement = block.getByPosition(arguments[2]).column;
+
+        ColumnWithTypeAndName & column_result = block.getByPosition(result);
+
+        bool needle_const = column_needle->isColumnConst();
+        bool replacement_const = column_replacement->isColumnConst();
+
+        if (column_src->isColumnConst())
+        {
+            executeImplConstHaystack(
+                column_src,
+                column_needle,
+                column_replacement,
+                needle_const,
+                replacement_const,
+                column_result);
+        }
+        else if (needle_const && replacement_const)
+        {
+            executeImpl(column_src, column_needle, column_replacement, column_result);
+        }
+        else if (needle_const)
+        {
+            executeImplNonConstReplacement(column_src, column_needle, column_replacement, column_result);
+        }
+        else if (replacement_const)
+        {
+            executeImplNonConstNeedle(column_src, column_needle, column_replacement, column_result);
+        }
+        else
+        {
+            executeImplNonConstNeedleReplacement(column_src, column_needle, column_replacement, column_result);
+        }
+    }
+
+private:
+    void executeImpl(
+        const ColumnPtr & column_src,
+        const ColumnPtr & column_needle,
+        const ColumnPtr & column_replacement,
+        ColumnWithTypeAndName & column_result) const
+    {
+        const auto * c1_const = typeid_cast<const ColumnConst *>(column_needle.get());
+        const auto * c2_const = typeid_cast<const ColumnConst *>(column_replacement.get());
+        auto needle = c1_const->getValue<String>();
+        auto replacement = c2_const->getValue<String>();
+
+        if (const auto * col = checkAndGetColumn<ColumnString>(column_src.get()))
+        {
+            auto col_res = ColumnString::create();
+            Impl::vector(
+                col->getChars(),
+                col->getOffsets(),
+                needle,
+                replacement,
+                col_res->getChars(),
+                col_res->getOffsets());
+            column_result.column = std::move(col_res);
+        }
+        else if (const auto * col = checkAndGetColumn<ColumnFixedString>(column_src.get()))
+        {
+            auto col_res = ColumnString::create();
+            Impl::vectorFixed(
+                col->getChars(),
+                col->getN(),
+                needle,
+                replacement,
+                col_res->getChars(),
+                col_res->getOffsets());
+            column_result.column = std::move(col_res);
+        }
+        else
+            throw Exception(
+                "Illegal column " + column_src->getName() + " of first argument of function " + getName(),
+                ErrorCodes::ILLEGAL_COLUMN);
+    }
+
+    void executeImplConstHaystack(
+        const ColumnPtr & column_src,
+        const ColumnPtr & column_needle,
+        const ColumnPtr & column_replacement,
+        bool needle_const,
+        bool replacement_const,
+        ColumnWithTypeAndName & column_result) const
+    {
+        auto res_col = ColumnString::create();
+        res_col->reserve(column_src->size());
+
+        RUNTIME_CHECK_MSG(
+            !needle_const || !replacement_const,
+            "should not got here when all argments of replace are constant");
+
+        const auto * column_src_const = checkAndGetColumnConst<ColumnString>(column_src.get());
+        RUNTIME_CHECK(column_src_const);
+
+        using GatherUtils::ConstSource;
+        using GatherUtils::StringSource;
+        if (!needle_const && !replacement_const)
+        {
+            const auto * column_needle_string = checkAndGetColumn<ColumnString>(column_needle.get());
+            const auto * column_replacement_string = checkAndGetColumn<ColumnString>(column_replacement.get());
+            RUNTIME_CHECK(column_needle_string);
+            RUNTIME_CHECK(column_replacement_string);
+
+            GatherUtils::replace<Impl>(
+                ConstSource<StringSource>(*column_src_const),
+                StringSource(*column_needle_string),
+                StringSource(*column_replacement_string),
+                res_col);
+        }
+        else if (needle_const && !replacement_const)
+        {
+            const auto * column_needle_const = checkAndGetColumnConst<ColumnString>(column_needle.get());
+            const auto * column_replacement_string = checkAndGetColumn<ColumnString>(column_replacement.get());
+            RUNTIME_CHECK(column_needle_const);
+            RUNTIME_CHECK(column_replacement_string);
+
+            GatherUtils::replace<Impl>(
+                ConstSource<StringSource>(*column_src_const),
+                ConstSource<StringSource>(*column_needle_const),
+                StringSource(*column_replacement_string),
+                res_col);
+        }
+        else if (!needle_const && replacement_const)
+        {
+            const auto * column_needle_string = checkAndGetColumn<ColumnString>(column_needle.get());
+            const auto * column_replacement_const = checkAndGetColumnConst<ColumnString>(column_replacement.get());
+            RUNTIME_CHECK(column_needle_string);
+            RUNTIME_CHECK(column_replacement_const);
+
+            GatherUtils::replace<Impl>(
+                ConstSource<StringSource>(*column_src_const),
+                StringSource(*column_needle_string),
+                ConstSource<StringSource>(*column_replacement_const),
+                res_col);
+        }
+
+        column_result.column = std::move(res_col);
+    }
+
+    void executeImplNonConstNeedle(
+        const ColumnPtr & column_src,
+        const ColumnPtr & column_needle,
+        const ColumnPtr & column_replacement,
+        ColumnWithTypeAndName & column_result) const
+    {
+        if constexpr (Impl::support_non_const_needle)
+        {
+            const auto * col_needle = typeid_cast<const ColumnString *>(column_needle.get());
+            const auto * col_replacement_const = typeid_cast<const ColumnConst *>(column_replacement.get());
+            auto replacement = col_replacement_const->getValue<String>();
+
+            if (const auto * col = checkAndGetColumn<ColumnString>(column_src.get()))
+            {
+                auto col_res = ColumnString::create();
+                Impl::vectorNonConstNeedle(
+                    col->getChars(),
+                    col->getOffsets(),
+                    col_needle->getChars(),
+                    col_needle->getOffsets(),
+                    replacement,
+                    col_res->getChars(),
+                    col_res->getOffsets());
+                column_result.column = std::move(col_res);
+            }
+            else if (const auto * col = checkAndGetColumn<ColumnFixedString>(column_src.get()))
+            {
+                auto col_res = ColumnString::create();
+                Impl::vectorFixedNonConstNeedle(
+                    col->getChars(),
+                    col->getN(),
+                    col_needle->getChars(),
+                    col_needle->getOffsets(),
+                    replacement,
+                    col_res->getChars(),
+                    col_res->getOffsets());
+                column_result.column = std::move(col_res);
+            }
+            else
+                throw Exception(
+                    "Illegal column " + column_src->getName() + " of first argument of function " + getName(),
+                    ErrorCodes::ILLEGAL_COLUMN);
+        }
+        else
+        {
+            throw Exception("Argument at index 2 for function replace must be constant", ErrorCodes::ILLEGAL_COLUMN);
+        }
+    }
+
+    void executeImplNonConstReplacement(
+        const ColumnPtr & column_src,
+        const ColumnPtr & column_needle,
+        const ColumnPtr & column_replacement,
+        ColumnWithTypeAndName & column_result) const
+    {
+        if constexpr (Impl::support_non_const_replacement)
+        {
+            const auto * col_needle_const = typeid_cast<const ColumnConst *>(column_needle.get());
+            auto needle = col_needle_const->getValue<String>();
+            const auto * col_replacement = typeid_cast<const ColumnString *>(column_replacement.get());
+
+            if (const auto * col = checkAndGetColumn<ColumnString>(column_src.get()))
+            {
+                auto col_res = ColumnString::create();
+                Impl::vectorNonConstReplacement(
+                    col->getChars(),
+                    col->getOffsets(),
+                    needle,
+                    col_replacement->getChars(),
+                    col_replacement->getOffsets(),
+                    col_res->getChars(),
+                    col_res->getOffsets());
+                column_result.column = std::move(col_res);
+            }
+            else if (const auto * col = checkAndGetColumn<ColumnFixedString>(column_src.get()))
+            {
+                auto col_res = ColumnString::create();
+                Impl::vectorFixedNonConstReplacement(
+                    col->getChars(),
+                    col->getN(),
+                    needle,
+                    col_replacement->getChars(),
+                    col_replacement->getOffsets(),
+                    col_res->getChars(),
+                    col_res->getOffsets());
+                column_result.column = std::move(col_res);
+            }
+            else
+                throw Exception(
+                    "Illegal column " + column_src->getName() + " of first argument of function " + getName(),
+                    ErrorCodes::ILLEGAL_COLUMN);
+        }
+        else
+        {
+            throw Exception("Argument at index 3 for function replace must be constant", ErrorCodes::ILLEGAL_COLUMN);
+        }
+    }
+
+    void executeImplNonConstNeedleReplacement(
+        const ColumnPtr & column_src,
+        const ColumnPtr & column_needle,
+        const ColumnPtr & column_replacement,
+        ColumnWithTypeAndName & column_result) const
+    {
+        if constexpr (Impl::support_non_const_needle && Impl::support_non_const_replacement)
+        {
+            const auto * col_needle = typeid_cast<const ColumnString *>(column_needle.get());
+            const auto * col_replacement = typeid_cast<const ColumnString *>(column_replacement.get());
+
+            if (const auto * col = checkAndGetColumn<ColumnString>(column_src.get()))
+            {
+                auto col_res = ColumnString::create();
+                Impl::vectorNonConstNeedleReplacement(
+                    col->getChars(),
+                    col->getOffsets(),
+                    col_needle->getChars(),
+                    col_needle->getOffsets(),
+                    col_replacement->getChars(),
+                    col_replacement->getOffsets(),
+                    col_res->getChars(),
+                    col_res->getOffsets());
+                column_result.column = std::move(col_res);
+            }
+            else if (const auto * col = checkAndGetColumn<ColumnFixedString>(column_src.get()))
+            {
+                auto col_res = ColumnString::create();
+                Impl::vectorFixedNonConstNeedleReplacement(
+                    col->getChars(),
+                    col->getN(),
+                    col_needle->getChars(),
+                    col_needle->getOffsets(),
+                    col_replacement->getChars(),
+                    col_replacement->getOffsets(),
+                    col_res->getChars(),
+                    col_res->getOffsets());
+                column_result.column = std::move(col_res);
+            }
+            else
+                throw Exception(
+                    "Illegal column " + column_src->getName() + " of first argument of function " + getName(),
+                    ErrorCodes::ILLEGAL_COLUMN);
+        }
+        else
+        {
+            throw Exception(
+                "Argument at index 2 and 3 for function replace must be constant",
+                ErrorCodes::ILLEGAL_COLUMN);
+        }
+    }
+};
+} // namespace DB