Skip to content

Commit e2bf850

Browse files
vitlibarvitlibar
committed
Add many math functions
1 parent 529fe7f commit e2bf850

File tree

6 files changed

+788
-0
lines changed

6 files changed

+788
-0
lines changed

‎src/Storages/TimeSeries/PrometheusQueryToSQL/applyFunction.cpp‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,11 @@
22

33
#include <Common/Exception.h>
44
#include <Storages/TimeSeries/PrometheusQueryToSQL/applyDateTimeFunction.h>
5+
#include <Storages/TimeSeries/PrometheusQueryToSQL/applyMathSimpleFunction.h>
56
#include <Storages/TimeSeries/PrometheusQueryToSQL/applyFunctionOverRange.h>
67
#include <Storages/TimeSeries/PrometheusQueryToSQL/applyFunctionScalar.h>
78
#include <Storages/TimeSeries/PrometheusQueryToSQL/applyFunctionVector.h>
9+
#include <Storages/TimeSeries/PrometheusQueryToSQL/fromFunctionPi.h>
810
#include <Storages/TimeSeries/PrometheusQueryToSQL/fromFunctionTime.h>
911

1012

@@ -33,6 +35,12 @@ SQLQueryPiece applyFunction(const PQT::Function * function_node, std::vector<SQL
3335
if (isDateTimeFunction(function_name))
3436
return applyDateTimeFunction(function_node, std::move(arguments), context);
3537

38+
if (isMathSimpleFunction(function_name))
39+
return applyMathSimpleFunction(function_node, std::move(arguments), context);
40+
41+
if (isFunctionPi(function_name))
42+
return fromFunctionPi(function_node, std::move(arguments), context);
43+
3644
if (isFunctionOverRange(function_name))
3745
return applyFunctionOverRange(function_node, std::move(arguments), context);
3846

‎src/Storages/TimeSeries/PrometheusQueryToSQL/applyMathSimpleFunction.cpp‎

Lines changed: 298 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,298 @@
1+
#include <Storages/TimeSeries/PrometheusQueryToSQL/applyMathSimpleFunction.h>
2+
3+
#include <Parsers/ASTFunction.h>
4+
#include <Parsers/ASTIdentifier.h>
5+
#include <Parsers/ASTLiteral.h>
6+
#include <Storages/TimeSeries/PrometheusQueryToSQL/ConverterContext.h>
7+
#include <Storages/TimeSeries/PrometheusQueryToSQL/SelectQueryBuilder.h>
8+
#include <Storages/TimeSeries/PrometheusQueryToSQL/dropMetricName.h>
9+
#include <Storages/TimeSeries/timeSeriesTypesToAST.h>
10+
#include <boost/math/special_functions/sign.hpp>
11+
#include <numbers>
12+
13+
14+
namespace DB::ErrorCodes
15+
{
16+
extern const int CANNOT_EXECUTE_PROMQL_QUERY;
17+
}
18+
19+
20+
namespace DB::PrometheusQueryToSQL
21+
{
22+
23+
namespace
24+
{
25+
/// Checks if the types of the specified arguments are valid for a math function.
26+
void checkArgumentTypes(const PQT::Function * function_node, const std::vector<SQLQueryPiece> & arguments, const ConverterContext & context)
27+
{
28+
const auto & function_name = function_node->function_name;
29+
30+
if (arguments.size() != 1)
31+
{
32+
throw Exception(ErrorCodes::CANNOT_EXECUTE_PROMQL_QUERY,
33+
"Function '{}' expects {} arguments, but was called with {} arguments",
34+
function_name, 1, arguments.size());
35+
}
36+
37+
const auto & argument = arguments[0];
38+
39+
if (argument.type != ResultType::INSTANT_VECTOR)
40+
{
41+
throw Exception(ErrorCodes::CANNOT_EXECUTE_PROMQL_QUERY,
42+
"Function '{}' expects an argument of type {}, but expression {} has type {}",
43+
function_name, ResultType::INSTANT_VECTOR,
44+
getPromQLText(argument, context), argument.type);
45+
}
46+
}
47+
48+
using EvaluateWithConstArgumentFunc = Float64 (*)(Float64);
49+
50+
struct ImplInfo
51+
{
52+
std::string_view ch_function_name;
53+
EvaluateWithConstArgumentFunc evaluate_with_const_argument;
54+
};
55+
56+
const ImplInfo * getImplInfo(std::string_view function_name)
57+
{
58+
static const std::unordered_map<std::string_view, ImplInfo> impl_map = {
59+
{"abs",
60+
{
61+
"abs",
62+
[](Float64 x) -> Float64 { return fabs(x); },
63+
}},
64+
65+
{"sgn",
66+
{
67+
"sign",
68+
[](Float64 x) -> Float64 { return boost::math::sign(x); },
69+
}},
70+
71+
{"floor",
72+
{
73+
"floor",
74+
[](Float64 x) -> Float64 { return floor(x); },
75+
}},
76+
77+
{"ceil",
78+
{
79+
"ceil",
80+
[](Float64 x) -> Float64 { return ceil(x); },
81+
}},
82+
83+
{"sqrt",
84+
{
85+
"sqrt",
86+
[](Float64 x) -> Float64 { return sqrt(x); },
87+
}},
88+
89+
{"exp",
90+
{
91+
"exp",
92+
[](Float64 x) -> Float64 { return exp(x); },
93+
}},
94+
95+
{"ln",
96+
{
97+
"log",
98+
[](Float64 x) -> Float64 { return log(x); },
99+
}},
100+
101+
{"log2",
102+
{
103+
"log2",
104+
[](Float64 x) -> Float64 { return log2(x); },
105+
}},
106+
107+
{"log10",
108+
{
109+
"log10",
110+
[](Float64 x) -> Float64 { return log10(x); },
111+
}},
112+
113+
{"rad",
114+
{
115+
"radians",
116+
[](Float64 x) -> Float64 { return x * (std::numbers::pi / 180.0); },
117+
}},
118+
119+
{"deg",
120+
{
121+
"degrees",
122+
[](Float64 x) -> Float64 { return x * (180.0 / std::numbers::pi); },
123+
}},
124+
125+
{"sin",
126+
{
127+
"sin",
128+
[](Float64 x) -> Float64 { return sin(x); },
129+
}},
130+
131+
{"cos",
132+
{
133+
"cos",
134+
[](Float64 x) -> Float64 { return cos(x); },
135+
}},
136+
137+
{"tan",
138+
{
139+
"tan",
140+
[](Float64 x) -> Float64 { return tan(x); },
141+
}},
142+
143+
{"asin",
144+
{
145+
"asin",
146+
[](Float64 x) -> Float64 { return asin(x); },
147+
}},
148+
149+
{"acos",
150+
{
151+
"acos",
152+
[](Float64 x) -> Float64 { return acos(x); },
153+
}},
154+
155+
{"atan",
156+
{
157+
"atan",
158+
[](Float64 x) -> Float64 { return atan(x); },
159+
}},
160+
161+
{"sinh",
162+
{
163+
"sinh",
164+
[](Float64 x) -> Float64 { return sinh(x); },
165+
}},
166+
167+
{"cosh",
168+
{
169+
"cosh",
170+
[](Float64 x) -> Float64 { return cosh(x); },
171+
}},
172+
173+
{"tanh",
174+
{
175+
"tanh",
176+
[](Float64 x) -> Float64 { return tanh(x); },
177+
}},
178+
179+
{"asinh",
180+
{
181+
"asinh",
182+
[](Float64 x) -> Float64 { return asinh(x); },
183+
}},
184+
185+
{"acosh",
186+
{
187+
"acosh",
188+
[](Float64 x) -> Float64 { return acosh(x); },
189+
}},
190+
191+
{"atanh",
192+
{
193+
"atanh",
194+
[](Float64 x) -> Float64 { return atanh(x); },
195+
}},
196+
};
197+
198+
auto it = impl_map.find(function_name);
199+
if (it == impl_map.end())
200+
return nullptr;
201+
202+
return &it->second;
203+
}
204+
}
205+
206+
207+
bool isMathSimpleFunction(std::string_view function_name)
208+
{
209+
return getImplInfo(function_name) != nullptr;
210+
}
211+
212+
213+
SQLQueryPiece applyMathSimpleFunction(
214+
const PQT::Function * function_node, std::vector<SQLQueryPiece> && arguments, ConverterContext & context)
215+
{
216+
const auto & function_name = function_node->function_name;
217+
const auto * impl_info = getImplInfo(function_name);
218+
chassert(impl_info);
219+
220+
checkArgumentTypes(function_node, arguments, context);
221+
auto & argument = arguments[0];
222+
223+
auto res = argument;
224+
res.node = function_node;
225+
226+
switch (argument.store_method)
227+
{
228+
case StoreMethod::EMPTY:
229+
{
230+
return res;
231+
}
232+
233+
case StoreMethod::CONST_SCALAR:
234+
{
235+
res.scalar_value = (impl_info->evaluate_with_const_argument)(argument.scalar_value);
236+
return res;
237+
}
238+
239+
case StoreMethod::SINGLE_SCALAR:
240+
{
241+
/// SELECT f(value) AS value FROM <subquery>
242+
SelectQueryBuilder builder;
243+
244+
builder.select_list.push_back(makeASTFunction(impl_info->ch_function_name, make_intrusive<ASTIdentifier>(ColumnNames::Value)));
245+
builder.select_list.back()->setAlias(ColumnNames::Value);
246+
247+
context.subqueries.emplace_back(SQLSubquery{context.subqueries.size(), std::move(argument.select_query), SQLSubqueryType::TABLE});
248+
builder.from_table = context.subqueries.back().name;
249+
250+
res.select_query = builder.getSelectQuery();
251+
return res;
252+
}
253+
254+
case StoreMethod::SCALAR_GRID:
255+
case StoreMethod::VECTOR_GRID:
256+
{
257+
/// For scalar grid:
258+
/// SELECT arrayMap(x -> f(x), values) AS values
259+
/// FROM <scalar_grid>
260+
///
261+
/// For vector grid:
262+
/// SELECT group, arrayMap(x -> f(x), values) AS values
263+
/// FROM <vector_grid>
264+
SelectQueryBuilder builder;
265+
if (argument.store_method == StoreMethod::VECTOR_GRID)
266+
builder.select_list.push_back(make_intrusive<ASTIdentifier>(ColumnNames::Group));
267+
268+
builder.select_list.push_back(makeASTFunction(
269+
"arrayMap",
270+
makeASTFunction(
271+
"lambda",
272+
makeASTFunction("tuple", make_intrusive<ASTIdentifier>("x")),
273+
makeASTFunction(impl_info->ch_function_name, make_intrusive<ASTIdentifier>("x"))),
274+
make_intrusive<ASTIdentifier>(ColumnNames::Values)));
275+
276+
builder.select_list.back()->setAlias(ColumnNames::Values);
277+
278+
context.subqueries.emplace_back(SQLSubquery{context.subqueries.size(), std::move(argument.select_query), SQLSubqueryType::TABLE});
279+
builder.from_table = context.subqueries.back().name;
280+
281+
res.select_query = builder.getSelectQuery();
282+
283+
return dropMetricName(std::move(res), context);
284+
}
285+
286+
case StoreMethod::CONST_STRING:
287+
case StoreMethod::RAW_DATA:
288+
{
289+
/// Can't get in here because these store methods are incompatible with the allowed argument types
290+
/// (see checkArgumentTypes()).
291+
throwUnexpectedStoreMethod(argument, context);
292+
}
293+
}
294+
295+
UNREACHABLE();
296+
}
297+
298+
}

‎src/Storages/TimeSeries/PrometheusQueryToSQL/applyMathSimpleFunction.h‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
#pragma once
2+
3+
#include <Storages/TimeSeries/PrometheusQueryToSQL/SQLQueryPiece.h>
4+
5+
6+
namespace DB::PrometheusQueryToSQL
7+
{
8+
9+
struct ConverterContext;
10+
11+
/// Returns whether a specified string is the name of a math function.
12+
/// The function only considers simple functions with one vector argument transforming each value separately.
13+
/// Examples: abs(), floor(), sqrt(), sin(), etc.
14+
bool isMathSimpleFunction(std::string_view function_name);
15+
16+
/// Applies a math function.
17+
SQLQueryPiece applyMathSimpleFunction(const PQT::Function * function_node, std::vector<SQLQueryPiece> && arguments, ConverterContext & context);
18+
19+
}

‎src/Storages/TimeSeries/PrometheusQueryToSQL/fromFunctionPi.cpp‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
#include <Storages/TimeSeries/PrometheusQueryToSQL/fromFunctionPi.h>
2+
3+
#include <Common/Exception.h>
4+
#include <Storages/TimeSeries/PrometheusQueryToSQL/ConverterContext.h>
5+
#include <numbers>
6+
7+
8+
namespace DB::ErrorCodes
9+
{
10+
extern const int CANNOT_EXECUTE_PROMQL_QUERY;
11+
}
12+
13+
14+
namespace DB::PrometheusQueryToSQL
15+
{
16+
17+
SQLQueryPiece fromFunctionPi(const PQT::Function * function_node, std::vector<SQLQueryPiece> && arguments, ConverterContext & context)
18+
{
19+
const auto & function_name = function_node->function_name;
20+
chassert(isFunctionPi(function_name));
21+
22+
if (!arguments.empty())
23+
{
24+
throw Exception(ErrorCodes::CANNOT_EXECUTE_PROMQL_QUERY,
25+
"Function '{}' expects no arguments, but was called with {} arguments",
26+
function_name, arguments.size());
27+
}
28+
29+
/// This implementation is similar to the implementation of function fromLiteral().
30+
auto node_range = context.node_range_getter.get(function_node);
31+
if (node_range.empty())
32+
return SQLQueryPiece{function_node, ResultType::SCALAR, StoreMethod::EMPTY};
33+
34+
SQLQueryPiece res{function_node, ResultType::SCALAR, StoreMethod::CONST_SCALAR};
35+
res.scalar_value = std::numbers::pi;
36+
res.start_time = node_range.start_time;
37+
res.end_time = node_range.end_time;
38+
res.step = node_range.step;
39+
40+
return res;
41+
}
42+
43+
}

0 commit comments

Comments
 (0)