be/src/udf/udf.h

// 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.
// This file is copied from
// https://github.com/apache/impala/blob/branch-2.9.0/be/src/udf/udf.h
// and modified by Doris

#pragma once

#include <gen_cpp/Types_types.h>

#include <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <vector>

#include "runtime/types.h"
#include "util/runtime_profile.h"
#include "vec/common/arena.h"

namespace doris {

struct ColumnPtrWrapper;
struct StringRef;
class RuntimeState;

namespace vectorized {
class IDataType;
using DataTypePtr = std::shared_ptr<const IDataType>;
} // namespace vectorized

// The FunctionContext is passed to every UDF/UDA and is the interface for the UDF to the
// rest of the system. It contains APIs to examine the system state, report errors
// and manage memory.
class FunctionContext {
public:
    enum FunctionStateScope {
        /// Indicates that the function state for this FunctionContext's UDF is shared across
        /// the plan fragment (a query is divided into multiple plan fragments, each of which
        /// is responsible for a part of the query execution). Within the plan fragment, there
        /// may be multiple instances of the UDF executing concurrently with multiple
        /// FunctionContexts sharing this state, meaning that the state must be
        /// thread-safe. The Prepare() function for the UDF may be called with this scope
        /// concurrently on a single host if the UDF will be evaluated in multiple plan
        /// fragments on that host. In general, read-only state that doesn't need to be
        /// recomputed for every UDF call should be fragment-local.
        /// TODO: not yet implemented
        FRAGMENT_LOCAL,

        /// Indicates that the function state is local to the execution thread. This state
        /// does not need to be thread-safe. However, this state will be initialized (via the
        /// Prepare() function) once for every execution thread, so fragment-local state
        /// should be used when possible for better performance. In general, inexpensive
        /// shared state that is written to by the UDF (e.g. scratch space) should be
        /// thread-local.
        THREAD_LOCAL,
    };

    static std::unique_ptr<doris::FunctionContext> create_context(
            RuntimeState* state, const vectorized::DataTypePtr& return_type,
            const std::vector<vectorized::DataTypePtr>& arg_types);

    /// Returns a new FunctionContext with the same constant args, fragment-local state, and
    /// debug flag as this FunctionContext. The caller is responsible for calling delete on
    /// it.
    std::unique_ptr<doris::FunctionContext> clone();

    void set_constant_cols(const std::vector<std::shared_ptr<doris::ColumnPtrWrapper>>& cols);

    RuntimeState* state() { return _state; }

    void set_dict_function(const TDictFunction& dict_function) { _dict_function = dict_function; }

    std::optional<TDictFunction>& dict_function() { return _dict_function; };
    bool check_overflow_for_decimal() const { return _check_overflow_for_decimal; }

    bool set_check_overflow_for_decimal(bool check_overflow_for_decimal) {
        return _check_overflow_for_decimal = check_overflow_for_decimal;
    }

    void set_string_as_jsonb_string(bool string_as_jsonb_string) {
        _string_as_jsonb_string = string_as_jsonb_string;
    }

    void set_jsonb_string_as_string(bool jsonb_string_as_string) {
        _jsonb_string_as_string = jsonb_string_as_string;
    }

    void set_udf_execute_timer(RuntimeProfile::Counter* udf_execute_timer) {
        _udf_execute_timer = udf_execute_timer;
    }

    RuntimeProfile::Counter* get_udf_execute_timer() { return _udf_execute_timer; }

    // Cast flag, when enable string_as_jsonb_string, string casting to jsonb will not parse string
    // instead just insert a string literal
    bool string_as_jsonb_string() const { return _string_as_jsonb_string; }

    // Cast flag, when enable jsonb_string_as_string, jsonb string casting to string will not parse string
    // instead just insert a string literal
    bool jsonb_string_as_string() const { return _jsonb_string_as_string; }

    // Sets an error for this UDF. If this is called, this will trigger the
    // query to fail.
    // Note: when you set error for the UDFs used in Data Load, you should
    // ensure the function return value is null.
    void set_error(const char* error_msg);

    // Adds a warning that is returned to the user. This can include things like
    // overflow or other recoverable error conditions.
    // Warnings are capped at a maximum number. Returns true if the warning was
    // added and false if it was ignored due to the cap.
    bool add_warning(const char* warning_msg);

    /// Methods for maintaining state across UDF/UDA function calls. SetFunctionState() can
    /// be used to store a pointer that can then be retrieved via GetFunctionState(). If
    /// GetFunctionState() is called when no pointer is set, it will return
    /// nullptr. SetFunctionState() does not take ownership of 'ptr'; it is up to the UDF/UDA
    /// to clean up any function state if necessary.
    void set_function_state(FunctionStateScope scope, std::shared_ptr<void> ptr);

    void* get_function_state(FunctionStateScope scope) const;

    // Returns the return type information of this function. For UDAs, this is the final
    // return type of the UDA (e.g., the type returned by the finalize function).
    const vectorized::DataTypePtr get_return_type() const;

    // Returns the number of arguments to this function (not including the FunctionContext*
    // argument).
    int get_num_args() const;

    // Returns the type information for the arg_idx-th argument (0-indexed, not including
    // the FunctionContext* argument). Returns nullptr if arg_idx is invalid.
    const vectorized::DataTypePtr get_arg_type(int arg_idx) const;

    // Returns true if the arg_idx-th input argument (0 indexed, not including the
    // FunctionContext* argument) is a constant (e.g. 5, "string", 1 + 1).
    bool is_col_constant(int arg_idx) const;

    // Returns a pointer to the value of the arg_idx-th input argument (0 indexed, not
    // including the FunctionContext* argument). Returns nullptr if the argument is not
    // constant. This function can be used to obtain user-specified constants in a UDF's
    // Init() or Close() functions.
    doris::ColumnPtrWrapper* get_constant_col(int arg_idx) const;

    // Creates a StringRef, which memory is available when this function context is used next time
    StringRef create_temp_string_val(int64_t len);

    ~FunctionContext() = default;

    vectorized::Arena& get_arena() { return arena; }

private:
    FunctionContext() = default;

    // Disable copy ctor and assignment operator
    FunctionContext(const FunctionContext& other);

    FunctionContext& operator=(const FunctionContext& other);

    // We use the query's runtime state to report errors and warnings. nullptr for test
    // contexts.
    RuntimeState* _state = nullptr;

    // Empty if there's no error
    std::string _error_msg;

    // The number of warnings reported.
    int64_t _num_warnings;

    /// The function state accessed via FunctionContext::Get/SetFunctionState()
    std::shared_ptr<void> _thread_local_fn_state;
    std::shared_ptr<void> _fragment_local_fn_state;

    // Type descriptor for the return type of the function.
    vectorized::DataTypePtr _return_type;

    // Type descriptors for each argument of the function.
    std::vector<vectorized::DataTypePtr> _arg_types;

    std::vector<std::shared_ptr<doris::ColumnPtrWrapper>> _constant_cols;

    //udf execute timer
    RuntimeProfile::Counter* _udf_execute_timer = nullptr;
    bool _check_overflow_for_decimal = false;

    bool _string_as_jsonb_string = false;
    bool _jsonb_string_as_string = false;

    std::string _string_result;

    vectorized::Arena arena;

    std::optional<TDictFunction> _dict_function;
};

using doris::FunctionContext;
} // namespace doris