-
Notifications
You must be signed in to change notification settings - Fork 9k
Expand file tree
/
Copy patharrayTopK.cpp
More file actions
326 lines (277 loc) · 13.1 KB
/
Copy patharrayTopK.cpp
File metadata and controls
326 lines (277 loc) · 13.1 KB
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
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
#include <Functions/array/arrayTopK.h>
#include <Columns/ColumnArray.h>
#include <Columns/ColumnDecimal.h>
#include <Columns/ColumnFixedString.h>
#include <Columns/ColumnNullable.h>
#include <Columns/ColumnString.h>
#include <Columns/ColumnsDateTime.h>
#include <Columns/ColumnsNumber.h>
#include <Functions/FunctionFactory.h>
#include <Functions/castTypeToEither.h>
#include <base/sort.h>
namespace DB
{
namespace ErrorCodes
{
extern const int BAD_ARGUMENTS;
extern const int LOGICAL_ERROR;
}
namespace
{
/// Comparator specialized for a concrete column type so `compareAt` is devirtualized and can be inlined.
template <bool IsAscending, typename ColumnType>
struct Less
{
const ColumnType & column;
explicit Less(const IColumn & column_)
: column(assert_cast<const ColumnType &>(column_))
{
}
bool operator()(size_t lhs, size_t rhs) const
{
int c = column.compareAt(lhs, rhs, column, IsAscending ? 1 : -1);
if (c != 0)
return IsAscending ? c < 0 : c > 0;
/// Make the result deterministic: equal elements keep their original order.
return lhs < rhs;
}
};
/// Comparator fallback for column types not covered by the specialized dispatch — uses virtual `compareAt`.
template <bool IsAscending>
struct GenericLess
{
const IColumn & column;
explicit GenericLess(const IColumn & column_) : column(column_) {}
bool operator()(size_t lhs, size_t rhs) const
{
int c = column.compareAt(lhs, rhs, column, IsAscending ? 1 : -1);
if (c != 0)
return IsAscending ? c < 0 : c > 0;
/// Make the result deterministic: equal elements keep their original order.
return lhs < rhs;
}
};
/// Reads K[row] from the K column and rejects negatives (signed columns).
size_t readK(const IColumn & k_column, bool k_is_signed, size_t row, const char * function_name)
{
const UInt64 k = k_column.getUInt(row);
/// For a signed K column, a negative value is reinterpreted as a huge UInt64 with the high bit set;
/// detect that and throw.
if (k_is_signed && k > static_cast<UInt64>(std::numeric_limits<Int64>::max()))
throw Exception(
ErrorCodes::BAD_ARGUMENTS,
"Argument K of function {} must be non-negative, got {}",
function_name,
static_cast<Int64>(k));
return k;
}
/// K-selection pass for a single call, specialized at compile time on the comparator type.
/// Builds the result `ColumnArray`.
template <typename Comparator>
ColumnPtr applyComparator(
const IColumn & k_column,
bool k_is_signed,
const ColumnArray & source,
const IColumn & mapped,
const char * function_name)
{
const auto & offsets = source.getOffsets();
const size_t size = offsets.size();
const size_t nested_size = source.getData().size();
/// Peel `Nullable` from `mapped` so the specialized `Less<T>` matches the underlying type,
/// and cache the null maps for direct array access in the inner loop.
const auto * mapped_nullable = checkAndGetColumn<ColumnNullable>(&mapped);
const IColumn & mapped_data = mapped_nullable ? mapped_nullable->getNestedColumn() : mapped;
const NullMap * mapped_null_map = mapped_nullable ? &mapped_nullable->getNullMapData() : nullptr;
/// Same peel for the source data column; the non-null nested column is what we use to build
/// the result (matches the declared return type).
const auto * source_nullable = checkAndGetColumn<ColumnNullable>(&source.getData());
const IColumn & source_data = source_nullable ? source_nullable->getNestedColumn() : source.getData();
const NullMap * source_null_map = source_nullable ? &source_nullable->getNullMapData() : nullptr;
auto indexes_column = ColumnUInt64::create();
auto & indexes = indexes_column->getData();
/// Smaller reserve when K is constant.
if (isColumnConst(k_column))
indexes.reserve(std::min(readK(k_column, k_is_signed, 0, function_name) * size, nested_size));
else
indexes.reserve(nested_size);
auto result_offsets_column = ColumnArray::ColumnOffsets::create(size);
auto & result_offsets = result_offsets_column->getData();
IColumn::Permutation row_indices;
ColumnArray::Offset current_offset = 0;
ColumnArray::Offset result_offset = 0;
Comparator cmp(mapped_data);
for (size_t i = 0; i < size; ++i)
{
const auto next_offset = offsets[i];
const size_t k = readK(k_column, k_is_signed, i, function_name);
if (!k)
{
result_offsets[i] = result_offset;
current_offset = next_offset;
continue;
}
row_indices.clear();
row_indices.reserve(next_offset - current_offset);
for (size_t j = current_offset; j < next_offset; ++j)
{
if (mapped_null_map && (*mapped_null_map)[j])
continue;
if (source_null_map && (*source_null_map)[j])
continue;
row_indices.push_back(j);
}
const size_t take = std::min(k, row_indices.size());
if (take == row_indices.size())
::sort(row_indices.begin(), row_indices.end(), cmp);
else
::partial_sort(row_indices.begin(), row_indices.begin() + take, row_indices.end(), cmp);
for (size_t j = 0; j < take; ++j)
indexes.push_back(row_indices[j]);
result_offset += take;
result_offsets[i] = result_offset;
current_offset = next_offset;
}
return ColumnArray::create(source_data.index(*indexes_column, 0), std::move(result_offsets_column));
}
/// Iterates the specialized-column type list, falling back to `GenericLess`.
template <bool IsAscending>
ColumnPtr dispatchByColumn(
const IColumn & k_column,
bool k_is_signed,
const ColumnArray & source,
const IColumn & mapped,
const char * function_name)
{
/// A constant lambda (e.g. `(x) -> NULL`) gives `mapped` as `ColumnConst(Nullable(...))`.
/// Materialize it so the `Nullable` peel below and the null-map access in `applyComparator` work uniformly.
auto mapped_full = mapped.convertToFullColumnIfConst();
/// To match a specialized `Less<T>` we look at the type underneath any `Nullable` wrapper.
const IColumn * mapped_data = mapped_full.get();
if (const auto * mapped_nullable = checkAndGetColumn<ColumnNullable>(mapped_full.get()))
mapped_data = &mapped_nullable->getNestedColumn();
/// Try using specialized comparator Less<>.
ColumnPtr result;
bool dispatched = castTypeToEither<
ColumnUInt8,
ColumnUInt16,
ColumnUInt32,
ColumnUInt64,
ColumnInt8,
ColumnInt16,
ColumnInt32,
ColumnInt64,
ColumnFloat32,
ColumnFloat64,
ColumnDateTime64,
ColumnDecimal<Decimal32>,
ColumnDecimal<Decimal64>,
ColumnDecimal<Decimal128>,
ColumnDecimal<Decimal256>,
ColumnString,
ColumnFixedString>(
mapped_data,
[&](const auto & column)
{
using ColumnT = std::decay_t<decltype(column)>;
result = applyComparator<Less<IsAscending, ColumnT>>(k_column, k_is_signed, source, *mapped_full, function_name);
return true;
});
if (dispatched)
return result;
/// Fall back to GenericLess<>.
return applyComparator<GenericLess<IsAscending>>(k_column, k_is_signed, source, *mapped_full, function_name);
}
}
template <bool IsAscending>
ColumnPtr ArrayTopKImpl<IsAscending>::execute(
const ColumnArray & array,
ColumnPtr mapped,
const ColumnWithTypeAndName * fixed_arguments)
{
const char * function_name = IsAscending ? "arrayBottomK" : "arrayTopK";
if (!fixed_arguments)
throw Exception(
ErrorCodes::LOGICAL_ERROR,
"Expected fixed arguments to get K for {}",
function_name);
const IColumn & k_column = *fixed_arguments[0].column;
const bool k_is_signed = isNativeInt(*fixed_arguments[0].type);
return dispatchByColumn<IsAscending>(k_column, k_is_signed, array, *mapped, function_name);
}
REGISTER_FUNCTION(ArrayTopK)
{
FunctionDocumentation::Description description = R"(
Returns an array of the K largest elements of the input array, sorted in descending order.
If a lambda function `f` is specified, elements are compared by the result of `f` applied to each element.
If `f` accepts multiple arguments, additional arrays are passed to `arrayTopK`, their elements
correspond to the arguments of `f`.
`NULL` values are skipped and do not appear in the result. The result size is at most `K`
and may be smaller when the input array contains fewer non-null elements than `K`.
The element type of the result is the non-nullable counterpart of the input element type.
`arrayTopK` is a [higher-order function](/reference/functions/regular-functions/overview#higher-order-functions).
See also `arrayBottomK`, which returns the K smallest elements instead.
)";
FunctionDocumentation::Syntax syntax = "arrayTopK([f,] K, arr [, arr1, ... ,arrN])";
FunctionDocumentation::Arguments arguments = {
{"f(arr[, arr1, ... ,arrN])", "Optional. A lambda function to compute the sort key for each element.", {"Lambda function"}},
{"K", "The number of largest elements to return.", {"(U)Int8/16/32/64"}},
{"arr", "An array.", {"Array(T)"}},
{"arr1, ... ,arrN", "N additional arrays, in the case when `f` accepts multiple arguments.", {"Array(T)"}}
};
FunctionDocumentation::ReturnedValue returned_value = {R"(
Returns up to `K` elements of `arr` with the largest values (or largest lambda results), sorted in descending order.
Nulls are skipped. The returned array has element type `T` even when the input has type `Nullable(T)`.
)"};
FunctionDocumentation::Examples examples = {
{"simple_int", "SELECT arrayTopK(3, [1, 5, 2, 7, 3])", "[7,5,3]"},
{"skip_nulls", "SELECT arrayTopK(3, [1, NULL, 5, 2, NULL, 7])", "[7,5,2]"},
{"fewer_than_k", "SELECT arrayTopK(5, [1, NULL, 2])", "[2,1]"},
{"lambda_simple", "SELECT arrayTopK((x) -> -x, 2, [5, 9, 1, 3])", "[1,3]"},
{"lambda_multi", "SELECT arrayTopK((x, y) -> y, 2, ['a', 'b', 'c'], [3, 1, 2])", "['a','c']"}
};
FunctionDocumentation::IntroducedIn introduced_in = {26, 6};
FunctionDocumentation::Category category = FunctionDocumentation::Category::Array;
FunctionDocumentation documentation = {description, syntax, arguments, {}, returned_value, examples, introduced_in, category};
factory.registerFunction<FunctionArrayTopK>(documentation);
}
REGISTER_FUNCTION(ArrayBottomK)
{
FunctionDocumentation::Description description = R"(
Returns an array of the K smallest elements of the input array, sorted in ascending order.
If a lambda function `f` is specified, elements are compared by the result of `f` applied to each element.
If `f` accepts multiple arguments, additional arrays are passed to `arrayBottomK`, their elements
correspond to the arguments of `f`.
`NULL` values are skipped and do not appear in the result. The result size is at most `K`
and may be smaller when the input array contains fewer non-null elements than `K`.
The element type of the result is the non-nullable counterpart of the input element type.
`arrayBottomK` is a [higher-order function](/reference/functions/regular-functions/overview#higher-order-functions).
See also:
- `arrayTopK`, which returns the K largest elements instead.
- `arrayPartialSort`, which produces the same K elements at positions `[1..K]` but
also keeps the remaining elements in unspecified order, and does not skip nulls.
)";
FunctionDocumentation::Syntax syntax = "arrayBottomK([f,] K, arr [, arr1, ... ,arrN])";
FunctionDocumentation::Arguments arguments = {
{"f(arr[, arr1, ... ,arrN])", "Optional. A lambda function to compute the sort key for each element.", {"Lambda function"}},
{"K", "The number of smallest elements to return.", {"(U)Int8/16/32/64"}},
{"arr", "An array.", {"Array(T)"}},
{"arr1, ... ,arrN", "N additional arrays, in the case when `f` accepts multiple arguments.", {"Array(T)"}}
};
FunctionDocumentation::ReturnedValue returned_value = {R"(
Returns up to `K` elements of `arr` with the smallest values (or smallest lambda results), sorted in ascending order.
Nulls are skipped. The returned array has element type `T` even when the input has type `Nullable(T)`.
)"};
FunctionDocumentation::Examples examples = {
{"simple_int", "SELECT arrayBottomK(3, [1, 5, 2, 7, 3])", "[1,2,3]"},
{"skip_nulls", "SELECT arrayBottomK(3, [1, NULL, 5, 2, NULL, 7])", "[1,2,5]"},
{"fewer_than_k", "SELECT arrayBottomK(5, [1, NULL, 2])", "[1,2]"},
{"lambda_simple", "SELECT arrayBottomK((x) -> -x, 2, [5, 9, 1, 3])", "[9,5]"},
{"lambda_multi", "SELECT arrayBottomK((x, y) -> y, 2, ['a', 'b', 'c'], [3, 1, 2])", "['b','c']"}
};
FunctionDocumentation::IntroducedIn introduced_in = {26, 6};
FunctionDocumentation::Category category = FunctionDocumentation::Category::Array;
FunctionDocumentation documentation = {description, syntax, arguments, {}, returned_value, examples, introduced_in, category};
factory.registerFunction<FunctionArrayBottomK>(documentation);
}
}