Merge pull request #51472 from bigo-sg/substring_index

Add function substring_index, as in spark or mysql

docs/en/sql-reference/functions/string-functions.md

@@ -575,6 +575,42 @@ Alias:
 
 Like `substring` but for Unicode code points. Assumes that the string contains valid UTF-8 encoded text. If this assumption is violated, no exception is thrown and the result is undefined.
 
+
+## substringIndex(s, delim, count)
+
+Returns the substring of `s` before `count` occurrences of the delimiter `delim`, as in Spark or MySQL.
+
+**Syntax**
+
+```sql
+substringIndex(s, delim, count)
+```
+Alias: `SUBSTRING_INDEX`
+
+
+**Arguments**
+
+- s: The string to extract substring from. [String](../../sql-reference/data-types/string.md).
+- delim: The character to split. [String](../../sql-reference/data-types/string.md).
+- count: The number of occurrences of the delimiter to count before extracting the substring. If count is positive, everything to the left of the final delimiter (counting from the left) is returned. If count is negative, everything to the right of the final delimiter (counting from the right) is returned. [UInt or Int](../data-types/int-uint.md)
+
+**Example**
+
+``` sql
+SELECT substringIndex('www.clickhouse.com', '.', 2)
+```
+
+Result:
+```
+┌─substringIndex('www.clickhouse.com', '.', 2)─┐
+│ www.clickhouse                               │
+└──────────────────────────────────────────────┘
+```
+
+## substringIndexUTF8(s, delim, count)
+
+Like `substringIndex` but for Unicode code points. Assumes that the string contains valid UTF-8 encoded text. If this assumption is violated, no exception is thrown and the result is undefined.
+
 ## appendTrailingCharIfAbsent
 
 Appends character `c` to string `s` if `s` is non-empty and does not end with character `c`.

src/Functions/substringIndex.cpp

@@ -0,0 +1,302 @@
+#include <Columns/ColumnConst.h>
+#include <Columns/ColumnString.h>
+#include <DataTypes/DataTypeString.h>
+#include <Functions/FunctionFactory.h>
+#include <Functions/FunctionHelpers.h>
+#include <Functions/IFunction.h>
+#include <Functions/PositionImpl.h>
+#include <Interpreters/Context_fwd.h>
+#include <base/find_symbols.h>
+#include <Common/UTF8Helpers.h>
+#include <Common/register_objects.h>
+
+namespace DB
+{
+
+namespace ErrorCodes
+{
+    extern const int ILLEGAL_COLUMN;
+    extern const int ILLEGAL_TYPE_OF_ARGUMENT;
+    extern const int BAD_ARGUMENTS;
+}
+
+namespace
+{
+
+    template <bool is_utf8>
+    class FunctionSubstringIndex : public IFunction
+    {
+    public:
+        static constexpr auto name = is_utf8 ? "substringIndexUTF8" : "substringIndex";
+
+
+        static FunctionPtr create(ContextPtr) { return std::make_shared<FunctionSubstringIndex>(); }
+
+        String getName() const override { return name; }
+
+        size_t getNumberOfArguments() const override { return 3; }
+
+        bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return true; }
+
+        bool useDefaultImplementationForConstants() const override { return true; }
+        ColumnNumbers getArgumentsThatAreAlwaysConstant() const override { return {1}; }
+
+        DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
+        {
+            if (!isString(arguments[0]))
+                throw Exception(
+                    ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
+                    "Illegal type {} of first argument of function {}, String expected",
+                    arguments[0]->getName(),
+                    getName());
+
+            if (!isString(arguments[1]))
+                throw Exception(
+                    ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
+                    "Illegal type {} of second argument of function {}, String expected",
+                    arguments[1]->getName(),
+                    getName());
+
+            if (!isNativeInteger(arguments[2]))
+                throw Exception(
+                    ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
+                    "Illegal type {} of third argument of function {}, Integer expected",
+                    arguments[2]->getName(),
+                    getName());
+
+            return std::make_shared<DataTypeString>();
+        }
+
+        ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t /*input_rows_count*/) const override
+        {
+            ColumnPtr column_string = arguments[0].column;
+            ColumnPtr column_delim = arguments[1].column;
+            ColumnPtr column_count = arguments[2].column;
+
+            const ColumnConst * column_delim_const = checkAndGetColumnConst<ColumnString>(column_delim.get());
+            if (!column_delim_const)
+                throw Exception(ErrorCodes::ILLEGAL_COLUMN, "Second argument to {} must be a constant String", getName());
+
+            String delim = column_delim_const->getValue<String>();
+            if constexpr (!is_utf8)
+            {
+                if (delim.size() != 1)
+                    throw Exception(ErrorCodes::BAD_ARGUMENTS, "Second argument to {} must be a single character", getName());
+            }
+            else
+            {
+                if (UTF8::countCodePoints(reinterpret_cast<const UInt8 *>(delim.data()), delim.size()) != 1)
+                    throw Exception(ErrorCodes::BAD_ARGUMENTS, "Second argument to {} must be a single UTF-8 character", getName());
+            }
+
+            auto column_res = ColumnString::create();
+            ColumnString::Chars & vec_res = column_res->getChars();
+            ColumnString::Offsets & offsets_res = column_res->getOffsets();
+
+            const ColumnConst * column_string_const = checkAndGetColumnConst<ColumnString>(column_string.get());
+            if (column_string_const)
+            {
+                String str = column_string_const->getValue<String>();
+                constantVector(str, delim, column_count.get(), vec_res, offsets_res);
+            }
+            else
+            {
+                const auto * col_str = checkAndGetColumn<ColumnString>(column_string.get());
+                if (!col_str)
+                    throw Exception(ErrorCodes::ILLEGAL_COLUMN, "First argument to {} must be a String", getName());
+
+                bool is_count_const = isColumnConst(*column_count);
+                if (is_count_const)
+                {
+                    Int64 count = column_count->getInt(0);
+                    vectorConstant(col_str, delim, count, vec_res, offsets_res);
+                }
+                else
+                    vectorVector(col_str, delim, column_count.get(), vec_res, offsets_res);
+            }
+            return column_res;
+        }
+
+    protected:
+        static void vectorVector(
+            const ColumnString * str_column,
+            const String & delim,
+            const IColumn * count_column,
+            ColumnString::Chars & res_data,
+            ColumnString::Offsets & res_offsets)
+        {
+            size_t rows = str_column->size();
+            res_data.reserve(str_column->getChars().size() / 2);
+            res_offsets.reserve(rows);
+
+            std::unique_ptr<PositionCaseSensitiveUTF8::SearcherInBigHaystack> searcher
+                = !is_utf8 ? nullptr : std::make_unique<PositionCaseSensitiveUTF8::SearcherInBigHaystack>(delim.data(), delim.size());
+
+            for (size_t i = 0; i < rows; ++i)
+            {
+                StringRef str_ref = str_column->getDataAt(i);
+                Int64 count = count_column->getInt(i);
+
+                StringRef res_ref;
+                if constexpr (!is_utf8)
+                    res_ref = substringIndex(str_ref, delim[0], count);
+                else
+                    res_ref = substringIndexUTF8(searcher.get(), str_ref, delim, count);
+
+                appendToResultColumn(res_ref, res_data, res_offsets);
+            }
+        }
+
+        static void vectorConstant(
+            const ColumnString * str_column,
+            const String & delim,
+            Int64 count,
+            ColumnString::Chars & res_data,
+            ColumnString::Offsets & res_offsets)
+        {
+            size_t rows = str_column->size();
+            res_data.reserve(str_column->getChars().size() / 2);
+            res_offsets.reserve(rows);
+
+            std::unique_ptr<PositionCaseSensitiveUTF8::SearcherInBigHaystack> searcher
+                = !is_utf8 ? nullptr : std::make_unique<PositionCaseSensitiveUTF8::SearcherInBigHaystack>(delim.data(), delim.size());
+
+            for (size_t i = 0; i < rows; ++i)
+            {
+                StringRef str_ref = str_column->getDataAt(i);
+
+                StringRef res_ref;
+                if constexpr (!is_utf8)
+                    res_ref = substringIndex(str_ref, delim[0], count);
+                else
+                    res_ref = substringIndexUTF8(searcher.get(), str_ref, delim, count);
+
+                appendToResultColumn(res_ref, res_data, res_offsets);
+            }
+        }
+
+        static void constantVector(
+            const String & str,
+            const String & delim,
+            const IColumn * count_column,
+            ColumnString::Chars & res_data,
+            ColumnString::Offsets & res_offsets)
+        {
+            size_t rows = count_column->size();
+            res_data.reserve(str.size() * rows / 2);
+            res_offsets.reserve(rows);
+
+            std::unique_ptr<PositionCaseSensitiveUTF8::SearcherInBigHaystack> searcher
+                = !is_utf8 ? nullptr : std::make_unique<PositionCaseSensitiveUTF8::SearcherInBigHaystack>(delim.data(), delim.size());
+
+            StringRef str_ref{str.data(), str.size()};
+            for (size_t i = 0; i < rows; ++i)
+            {
+                Int64 count = count_column->getInt(i);
+
+                StringRef res_ref;
+                if constexpr (!is_utf8)
+                    res_ref = substringIndex(str_ref, delim[0], count);
+                else
+                    res_ref = substringIndexUTF8(searcher.get(), str_ref, delim, count);
+
+                appendToResultColumn(res_ref, res_data, res_offsets);
+            }
+        }
+
+        static void appendToResultColumn(const StringRef & res_ref, ColumnString::Chars & res_data, ColumnString::Offsets & res_offsets)
+        {
+            size_t res_offset = res_data.size();
+            res_data.resize(res_offset + res_ref.size + 1);
+            memcpy(&res_data[res_offset], res_ref.data, res_ref.size);
+            res_offset += res_ref.size;
+            res_data[res_offset] = 0;
+            ++res_offset;
+
+            res_offsets.emplace_back(res_offset);
+        }
+
+        static StringRef substringIndexUTF8(
+            const PositionCaseSensitiveUTF8::SearcherInBigHaystack * searcher, const StringRef & str_ref, const String & delim, Int64 count)
+        {
+            if (count == 0)
+                return {str_ref.data, 0};
+
+            const auto * begin = reinterpret_cast<const UInt8 *>(str_ref.data);
+            const auto * end = reinterpret_cast<const UInt8 *>(str_ref.data + str_ref.size);
+            const auto * pos = begin;
+            if (count > 0)
+            {
+                Int64 i = 0;
+                while (i < count)
+                {
+                    pos = searcher->search(pos, end - pos);
+
+                    if (pos != end)
+                    {
+                        pos += delim.size();
+                        ++i;
+                    }
+                    else
+                        return str_ref;
+                }
+                return {begin, static_cast<size_t>(pos - begin - delim.size())};
+            }
+            else
+            {
+                Int64 total = 0;
+                while (pos < end && end != (pos = searcher->search(pos, end - pos)))
+                {
+                    pos += delim.size();
+                    ++total;
+                }
+
+                if (total + count < 0)
+                    return str_ref;
+
+                pos = begin;
+                Int64 i = 0;
+                Int64 count_from_left = total + 1 + count;
+                while (i < count_from_left && pos < end && end != (pos = searcher->search(pos, end - pos)))
+                {
+                    pos += delim.size();
+                    ++i;
+                }
+                return {pos, static_cast<size_t>(end - pos)};
+            }
+        }
+
+        static StringRef substringIndex(const StringRef & str_ref, char delim, Int64 count)
+        {
+            if (count == 0)
+                return {str_ref.data, 0};
+
+            const auto * pos = count > 0 ? str_ref.data : str_ref.data + str_ref.size - 1;
+            const auto * end = count > 0 ? str_ref.data + str_ref.size : str_ref.data - 1;
+            int d = count > 0 ? 1 : -1;
+
+            for (; count; pos += d)
+            {
+                if (pos == end)
+                    return str_ref;
+                if (*pos == delim)
+                    count -= d;
+            }
+            pos -= d;
+            return {
+                d > 0 ? str_ref.data : pos + 1, static_cast<size_t>(d > 0 ? pos - str_ref.data : str_ref.data + str_ref.size - pos - 1)};
+        }
+    };
+}
+
+
+REGISTER_FUNCTION(SubstringIndex)
+{
+    factory.registerFunction<FunctionSubstringIndex<false>>(); /// substringIndex
+    factory.registerFunction<FunctionSubstringIndex<true>>(); /// substringIndexUTF8
+
+    factory.registerAlias("SUBSTRING_INDEX", "substringIndex", FunctionFactory::CaseInsensitive);
+}
+
+
+}

tests/queries/0_stateless/02415_all_new_functions_must_be_documented.reference

@@ -667,6 +667,8 @@ sqrt
 startsWith
 subBitmap
 substring
+substringIndex
+substringIndexUTF8
 substringUTF8
 subtractDays
 subtractHours