1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
|
#include <Functions/FunctionHelpers.h>
#include <Functions/FunctionFactory.h>
#include <BridgeHelper/CatBoostLibraryBridgeHelper.h>
#include <BridgeHelper/IBridgeHelper.h>
#include <Columns/ColumnNullable.h>
#include <Columns/ColumnString.h>
#include <Columns/ColumnTuple.h>
#include <Columns/ColumnsNumber.h>
#include <Common/assert_cast.h>
#include <DataTypes/DataTypeNullable.h>
#include <DataTypes/DataTypeTuple.h>
#include <DataTypes/DataTypesNumber.h>
#include <Functions/IFunction.h>
#include <Interpreters/Context.h>
#include <Interpreters/Context_fwd.h>
namespace DB
{
namespace ErrorCodes
{
extern const int FILE_DOESNT_EXIST;
extern const int ILLEGAL_TYPE_OF_ARGUMENT;
extern const int TOO_FEW_ARGUMENTS_FOR_FUNCTION;
extern const int ILLEGAL_COLUMN;
}
/// Evaluate CatBoost model.
/// - Arguments: float features first, then categorical features.
/// - Result: Float64.
class FunctionCatBoostEvaluate final : public IFunction, WithContext
{
private:
mutable std::unique_ptr<CatBoostLibraryBridgeHelper> bridge_helper;
public:
static constexpr auto name = "catboostEvaluate";
static FunctionPtr create(ContextPtr context_) { return std::make_shared<FunctionCatBoostEvaluate>(context_); }
explicit FunctionCatBoostEvaluate(ContextPtr context_) : WithContext(context_) {}
String getName() const override { return name; }
bool isVariadic() const override { return true; }
bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return true; }
bool isDeterministic() const override { return false; }
bool useDefaultImplementationForNulls() const override { return false; }
size_t getNumberOfArguments() const override { return 0; }
void initBridge(const ColumnConst * name_col) const
{
String library_path = getContext()->getConfigRef().getString("catboost_lib_path");
if (!std::filesystem::exists(library_path))
throw Exception(ErrorCodes::FILE_DOESNT_EXIST, "Can't load library {}: file doesn't exist", library_path);
String model_path = name_col->getValue<String>();
if (!std::filesystem::exists(model_path))
throw Exception(ErrorCodes::FILE_DOESNT_EXIST, "Can't load model {}: file doesn't exist", model_path);
bridge_helper = std::make_unique<CatBoostLibraryBridgeHelper>(getContext(), model_path, library_path);
}
DataTypePtr getReturnTypeFromLibraryBridge() const
{
size_t tree_count = bridge_helper->getTreeCount();
auto type = std::make_shared<DataTypeFloat64>();
if (tree_count == 1)
return type;
DataTypes types(tree_count, type);
return std::make_shared<DataTypeTuple>(types);
}
DataTypePtr getReturnTypeImpl(const ColumnsWithTypeAndName & arguments) const override
{
if (arguments.size() < 2)
throw Exception(ErrorCodes::TOO_FEW_ARGUMENTS_FOR_FUNCTION, "Function {} expects at least 2 arguments", getName());
if (!isString(arguments[0].type))
throw Exception(
ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
"Illegal type {} of first argument of function {}, expected a string.", arguments[0].type->getName(), getName());
const auto * name_col = checkAndGetColumnConst<ColumnString>(arguments[0].column.get());
if (!name_col)
throw Exception(ErrorCodes::ILLEGAL_COLUMN, "First argument of function {} must be a constant string", getName());
initBridge(name_col);
auto type = getReturnTypeFromLibraryBridge();
bool has_nullable = false;
for (size_t i = 1; i < arguments.size(); ++i)
has_nullable = has_nullable || arguments[i].type->isNullable();
if (has_nullable)
{
if (const auto * tuple = typeid_cast<const DataTypeTuple *>(type.get()))
{
auto elements = tuple->getElements();
for (auto & element : elements)
element = makeNullable(element);
type = std::make_shared<DataTypeTuple>(elements);
}
else
type = makeNullable(type);
}
return type;
}
ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t) const override
{
const auto * name_col = checkAndGetColumnConst<ColumnString>(arguments[0].column.get());
if (!name_col)
throw Exception(ErrorCodes::ILLEGAL_COLUMN, "First argument of function {} must be a constant string", getName());
ColumnRawPtrs column_ptrs;
Columns materialized_columns;
ColumnPtr null_map;
ColumnsWithTypeAndName feature_arguments(arguments.begin() + 1, arguments.end());
for (auto & arg : feature_arguments)
{
if (auto full_column = arg.column->convertToFullColumnIfConst())
{
materialized_columns.push_back(full_column);
arg.column = full_column;
}
if (const auto * col_nullable = checkAndGetColumn<ColumnNullable>(&*arg.column))
{
if (!null_map)
null_map = col_nullable->getNullMapColumnPtr();
else
{
auto mut_null_map = IColumn::mutate(std::move(null_map));
NullMap & result_null_map = assert_cast<ColumnUInt8 &>(*mut_null_map).getData();
const NullMap & src_null_map = col_nullable->getNullMapColumn().getData();
for (size_t i = 0, size = result_null_map.size(); i < size; ++i)
if (src_null_map[i])
result_null_map[i] = 1;
null_map = std::move(mut_null_map);
}
arg.column = col_nullable->getNestedColumn().getPtr();
arg.type = static_cast<const DataTypeNullable &>(*arg.type).getNestedType();
}
}
auto res = bridge_helper->evaluate(feature_arguments);
if (null_map)
{
if (const auto * tuple = typeid_cast<const ColumnTuple *>(res.get()))
{
auto nested = tuple->getColumns();
for (auto & col : nested)
col = ColumnNullable::create(col, null_map);
res = ColumnTuple::create(nested);
}
else
res = ColumnNullable::create(res, null_map);
}
return res;
}
};
REGISTER_FUNCTION(CatBoostEvaluate)
{
factory.registerFunction<FunctionCatBoostEvaluate>();
}
}
|
