feat: Support elt pushdown

dbms/src/Flash/Coprocessor/DAGUtils.cpp

@@ -610,7 +610,7 @@ const std::unordered_map<tipb::ScalarFuncSig, String> 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"},

dbms/src/Functions/FunctionHelpers.h

@@ -55,14 +55,24 @@ bool checkColumn(const IColumn * column)
     return checkAndGetColumn<Type>(column);
 }
 
+template <typename Type>
+const Type * checkAndGetNestedColumn(const IColumn * column)
+{
+    if (!column || !column->isColumnNullable())
+        return {};
+
+    const auto * data_column = &static_cast<const ColumnNullable *>(column)->getNestedColumn();
+
+    return checkAndGetColumn<Type>(data_column);
+}
 
 template <typename Type>
 const ColumnConst * checkAndGetColumnConst(const IColumn * column, bool maybe_nullable_column = false)
 {
     if (!column || !column->isColumnConst())
         return {};
 
-    const ColumnConst * res = static_cast<const ColumnConst *>(column);
+    const auto * res = static_cast<const ColumnConst *>(column);
 
     const auto * data_column = &res->getDataColumn();
     if (maybe_nullable_column && data_column->isColumnNullable())

dbms/src/Functions/FunctionsString.cpp

@@ -29,15 +29,11 @@
 #include <Functions/GatherUtils/GatherUtils.h>
 #include <Functions/StringUtil.h>
 #include <Functions/castTypeToEither.h>
-#include <IO/WriteHelpers.h>
 #include <Interpreters/Context.h>
 #include <fmt/core.h>
-#include <fmt/format.h>
-#include <fmt/printf.h>
 
 #include <boost/algorithm/string/predicate.hpp>
 #include <ext/range.h>
-#include <thread>
 
 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<FunctionElt>();
+    }
+
+    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<DataTypeString>());
+    }
+
+    void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result) const override
+    {
+        if (executeElt<UInt8>(block, arguments, result)
+            || executeElt<UInt16>(block, arguments, result)
+            || executeElt<UInt32>(block, arguments, result)
+            || executeElt<UInt64>(block, arguments, result)
+            || executeElt<Int8>(block, arguments, result)
+            || executeElt<Int16>(block, arguments, result)
+            || executeElt<Int32>(block, arguments, result)
+            || executeElt<Int64>(block, arguments, result))
+        {
+            return;
+        }
+        else
+        {
+            throw Exception(fmt::format("Illegal argument of function {}", getName()), ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
+        }
+    }
 
 private:
+    using NullMapMutablePtr = COWPtrHelper<DB::ColumnVectorHelper, DB::ColumnVector<unsigned char>>::MutablePtr;
+
+    template <typename IntType>
+    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<ColumnVector<IntType>>(col_arg0, true))
+        {
+            return constColumn<IntType>(col, block, arguments, result);
+        }
+        else
+        {
+            return vectorColumn<IntType>(col_arg0, block, arguments, result);
+        }
+    }
+
+    static void fillResultColumnNull(ColumnPtr & dst, size_t nrow)
+    {
+        dst = DataTypeNullable(std::make_shared<DataTypeString>()).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<ColumnString>(src.get(), true)->getDataColumnPtr().get())
+            : src.get();
+
+        const auto * col_str = col_nullable_str->isColumnNullable()
+            ? checkAndGetNestedColumn<ColumnString>(col_nullable_str)
+            : checkAndGetColumn<ColumnString>(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 <typename IntType>
+    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<ColumnVector<IntType>>(col->getDataColumnPtr().get())->getInt(0)
+            : col->getInt(0);
+
+        if (arg0 < 1 || arg0 >= static_cast<Int64>(arguments.size()))
+        {
+            fillResultColumnNull(block.getByPosition(result).column, nrow);
+        }
+        else
+        {
+            fillResultColumnFromOther(block.getByPosition(result).column, block.getByPosition(arg0).column);
+        }
+        return true;
+    }
+
+    template <typename IntType>
+    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<ColumnVector<IntType>>(col)
+            : checkAndGetColumn<ColumnVector<IntType>>(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<Int64>(arg0) >= static_cast<Int64>(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<FunctionHexInt>();
     factory.registerFunction<FunctionRepeat>();
     factory.registerFunction<FunctionBin>();
+    factory.registerFunction<FunctionElt>();
 }
 } // namespace DB