be/src/exprs/function/function_levenshtein.cpp

// 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 <algorithm>
#include <vector>

#include "common/status.h"
#include "core/data_type/data_type_number.h"
#include "core/string_ref.h"
#include "exprs/function/function_totype.h"
#include "exprs/function/simple_function_factory.h"
#include "util/simd/vstring_function.h"

namespace doris {

struct NameLevenshtein {
    static constexpr auto name = "levenshtein";
};

template <typename LeftDataType, typename RightDataType>
struct LevenshteinImpl {
    using ResultDataType = DataTypeInt32;
    using ResultPaddedPODArray = PaddedPODArray<Int32>;

    static Status vector_vector(const ColumnString::Chars& ldata,
                                const ColumnString::Offsets& loffsets,
                                const ColumnString::Chars& rdata,
                                const ColumnString::Offsets& roffsets, ResultPaddedPODArray& res) {
        DCHECK_EQ(loffsets.size(), roffsets.size());

        const size_t size = loffsets.size();
        res.resize(size);
        std::vector<size_t> left_offsets;
        std::vector<size_t> right_offsets;
        for (size_t i = 0; i < size; ++i) {
            res[i] = levenshtein_distance(string_ref_at(ldata, loffsets, i),
                                          string_ref_at(rdata, roffsets, i), left_offsets,
                                          right_offsets);
        }
        return Status::OK();
    }

    static Status vector_scalar(const ColumnString::Chars& ldata,
                                const ColumnString::Offsets& loffsets, const StringRef& rdata,
                                ResultPaddedPODArray& res) {
        const size_t size = loffsets.size();
        res.resize(size);
        const auto right = rdata.trim_tail_padding_zero();
        const bool right_ascii = simd::VStringFunctions::is_ascii(right);
        std::vector<size_t> right_offsets;
        simd::VStringFunctions::get_utf8_char_offsets(right, right_offsets);
        std::vector<size_t> left_offsets;
        for (size_t i = 0; i < size; ++i) {
            res[i] = levenshtein_distance_with_right_offsets(string_ref_at(ldata, loffsets, i),
                                                             left_offsets, right, right_offsets,
                                                             right_ascii);
        }
        return Status::OK();
    }

    static Status scalar_vector(const StringRef& ldata, const ColumnString::Chars& rdata,
                                const ColumnString::Offsets& roffsets, ResultPaddedPODArray& res) {
        const size_t size = roffsets.size();
        res.resize(size);
        const auto left = ldata.trim_tail_padding_zero();
        const bool left_ascii = simd::VStringFunctions::is_ascii(left);
        std::vector<size_t> left_offsets;
        simd::VStringFunctions::get_utf8_char_offsets(left, left_offsets);
        std::vector<size_t> right_offsets;
        for (size_t i = 0; i < size; ++i) {
            res[i] = levenshtein_distance_with_left_offsets(left, left_offsets, left_ascii,
                                                            string_ref_at(rdata, roffsets, i),
                                                            right_offsets);
        }
        return Status::OK();
    }

private:
    static StringRef string_ref_at(const ColumnString::Chars& data,
                                   const ColumnString::Offsets& offsets, size_t i) {
        DCHECK_LT(i, offsets.size());
        const auto idx = static_cast<ssize_t>(i);
        return StringRef(data.data() + offsets[idx - 1], offsets[idx] - offsets[idx - 1])
                .trim_tail_padding_zero();
    }

    static Int32 levenshtein_distance_utf8(const StringRef& left,
                                           const std::vector<size_t>& left_offsets,
                                           const StringRef& right,
                                           const std::vector<size_t>& right_offsets) {
        const StringRef* left_ref = &left;
        const StringRef* right_ref = &right;
        const std::vector<size_t>* left_offsets_ref = &left_offsets;
        const std::vector<size_t>* right_offsets_ref = &right_offsets;
        if (right_offsets_ref->size() > left_offsets_ref->size()) {
            std::swap(left_offsets_ref, right_offsets_ref);
            std::swap(left_ref, right_ref);
        }

        const size_t m = left_offsets_ref->size();
        const size_t n = right_offsets_ref->size();

        std::vector<Int32> prev(n + 1);
        std::vector<Int32> curr(n + 1);
        for (size_t j = 0; j <= n; ++j) {
            prev[j] = static_cast<Int32>(j);
        }

        for (size_t i = 1; i <= m; ++i) {
            curr[0] = static_cast<Int32>(i);
            const size_t left_off = (*left_offsets_ref)[i - 1];
            const size_t left_next = i < m ? (*left_offsets_ref)[i] : left_ref->size;

            for (size_t j = 1; j <= n; ++j) {
                const size_t right_off = (*right_offsets_ref)[j - 1];
                const size_t right_next = j < n ? (*right_offsets_ref)[j] : right_ref->size;

                const Int32 cost =
                        simd::VStringFunctions::utf8_char_equal(*left_ref, left_off, left_next,
                                                                *right_ref, right_off, right_next)
                                ? 0
                                : 1;

                const Int32 insert_cost = curr[j - 1] + 1;
                const Int32 delete_cost = prev[j] + 1;
                const Int32 replace_cost = prev[j - 1] + cost;
                curr[j] = std::min(std::min(insert_cost, delete_cost), replace_cost);
            }
            std::swap(prev, curr);
        }

        return prev[n];
    }

    static Int32 levenshtein_distance_ascii(const StringRef& left, const StringRef& right) {
        const StringRef* left_ref = &left;
        const StringRef* right_ref = &right;
        size_t m = left.size;
        size_t n = right.size;

        if (n > m) {
            std::swap(left_ref, right_ref);
            std::swap(m, n);
        }

        std::vector<Int32> prev(n + 1);
        std::vector<Int32> curr(n + 1);
        for (size_t j = 0; j <= n; ++j) {
            prev[j] = static_cast<Int32>(j);
        }

        for (size_t i = 1; i <= m; ++i) {
            curr[0] = static_cast<Int32>(i);
            const char left_char = left_ref->data[i - 1];

            for (size_t j = 1; j <= n; ++j) {
                const Int32 cost = left_char == right_ref->data[j - 1] ? 0 : 1;
                const Int32 insert_cost = curr[j - 1] + 1;
                const Int32 delete_cost = prev[j] + 1;
                const Int32 replace_cost = prev[j - 1] + cost;
                curr[j] = std::min(std::min(insert_cost, delete_cost), replace_cost);
            }
            std::swap(prev, curr);
        }

        return prev[n];
    }

    static Int32 levenshtein_distance(const StringRef& left, const StringRef& right,
                                      std::vector<size_t>& left_offsets,
                                      std::vector<size_t>& right_offsets) {
        const bool left_ascii = simd::VStringFunctions::is_ascii(left);
        const bool right_ascii = simd::VStringFunctions::is_ascii(right);
        if (left_ascii && right_ascii) {
            return levenshtein_distance_ascii(left, right);
        }

        if (left.size == 0) {
            return static_cast<Int32>(simd::VStringFunctions::get_char_len(right.data, right.size));
        }
        if (right.size == 0) {
            return static_cast<Int32>(simd::VStringFunctions::get_char_len(left.data, left.size));
        }

        simd::VStringFunctions::get_utf8_char_offsets(left, left_offsets);
        simd::VStringFunctions::get_utf8_char_offsets(right, right_offsets);
        return levenshtein_distance_utf8(left, left_offsets, right, right_offsets);
    }

    static Int32 levenshtein_distance_with_right_offsets(const StringRef& left,
                                                         std::vector<size_t>& left_offsets,
                                                         const StringRef& right,
                                                         const std::vector<size_t>& right_offsets,
                                                         bool right_ascii) {
        const bool left_ascii = simd::VStringFunctions::is_ascii(left);
        if (left_ascii && right_ascii) {
            return levenshtein_distance_ascii(left, right);
        }

        if (left.size == 0) {
            return static_cast<Int32>(right_offsets.size());
        }
        if (right.size == 0) {
            return left_ascii ? static_cast<Int32>(left.size)
                              : static_cast<Int32>(
                                        simd::VStringFunctions::get_char_len(left.data, left.size));
        }

        simd::VStringFunctions::get_utf8_char_offsets(left, left_offsets);
        return levenshtein_distance_utf8(left, left_offsets, right, right_offsets);
    }

    static Int32 levenshtein_distance_with_left_offsets(const StringRef& left,
                                                        const std::vector<size_t>& left_offsets,
                                                        bool left_ascii, const StringRef& right,
                                                        std::vector<size_t>& right_offsets) {
        const bool right_ascii = simd::VStringFunctions::is_ascii(right);
        if (left_ascii && right_ascii) {
            return levenshtein_distance_ascii(left, right);
        }

        if (left.size == 0) {
            return static_cast<Int32>(
                    right_ascii ? right.size
                                : simd::VStringFunctions::get_char_len(right.data, right.size));
        }
        if (right.size == 0) {
            return static_cast<Int32>(left_offsets.size());
        }

        simd::VStringFunctions::get_utf8_char_offsets(right, right_offsets);
        return levenshtein_distance_utf8(left, left_offsets, right, right_offsets);
    }

    static Int32 levenshtein_distance(const StringRef& left, const StringRef& right) {
        std::vector<size_t> left_offsets;
        std::vector<size_t> right_offsets;
        return levenshtein_distance(left, right, left_offsets, right_offsets);
    }
};

using FunctionLevenshtein =
        FunctionBinaryToType<DataTypeString, DataTypeString, LevenshteinImpl, NameLevenshtein>;

void register_function_levenshtein(SimpleFunctionFactory& factory) {
    factory.register_function<FunctionLevenshtein>();
}

} // namespace doris