Coverage Report

Created: 2026-07-17 17:00

next uncovered line (L), next uncovered region (R), next uncovered branch (B)
be/src/exprs/lambda_function/varray_map_function.cpp
Line
Count
Source
1
// Licensed to the Apache Software Foundation (ASF) under one
2
// or more contributor license agreements.  See the NOTICE file
3
// distributed with this work for additional information
4
// regarding copyright ownership.  The ASF licenses this file
5
// to you under the Apache License, Version 2.0 (the
6
// "License"); you may not use this file except in compliance
7
// with the License.  You may obtain a copy of the License at
8
//
9
//   http://www.apache.org/licenses/LICENSE-2.0
10
//
11
// Unless required by applicable law or agreed to in writing,
12
// software distributed under the License is distributed on an
13
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14
// KIND, either express or implied.  See the License for the
15
// specific language governing permissions and limitations
16
// under the License.
17
18
#include <algorithm>
19
#include <memory>
20
#include <set>
21
#include <string>
22
#include <vector>
23
24
#include "common/check.h"
25
#include "common/status.h"
26
#include "core/assert_cast.h"
27
#include "core/block/block.h"
28
#include "core/block/column_numbers.h"
29
#include "core/block/column_with_type_and_name.h"
30
#include "core/block/columns_with_type_and_name.h"
31
#include "core/column/column.h"
32
#include "core/column/column_array.h"
33
#include "core/column/column_nothing.h"
34
#include "core/column/column_nullable.h"
35
#include "core/column/column_vector.h"
36
#include "core/data_type/data_type.h"
37
#include "core/data_type/data_type_array.h"
38
#include "core/data_type/data_type_nullable.h"
39
#include "core/data_type/data_type_number.h"
40
#include "exec/common/util.hpp"
41
#include "exprs/aggregate/aggregate_function.h"
42
#include "exprs/lambda_function/lambda_execution_context.h"
43
#include "exprs/lambda_function/lambda_function.h"
44
#include "exprs/lambda_function/lambda_function_factory.h"
45
#include "exprs/vcolumn_ref.h"
46
#include "exprs/vexpr_context.h"
47
#include "exprs/vlambda_function_expr.h"
48
49
namespace doris {
50
51
// extend a block with all required parameters
52
struct LambdaArgs {
53
    // which line is extended to the original block
54
    int64_t current_row_idx = 0;
55
    // when a block is filled, the array may be truncated, recording where it was truncated
56
    int64_t current_offset_in_array = 0;
57
    // the beginning position of the array
58
    size_t array_start = 0;
59
    // the size of the array
60
    int64_t cur_size = 0;
61
    // offset of column array
62
    const ColumnArray::Offsets64* offsets_ptr = nullptr;
63
    // expend data of repeat times
64
    int current_repeat_times = 0;
65
    // whether the current row of the original block has been extended
66
    bool current_row_eos = false;
67
};
68
69
class ArrayMapFunction : public LambdaFunction {
70
    ENABLE_FACTORY_CREATOR(ArrayMapFunction);
71
72
public:
73
27
    ~ArrayMapFunction() override = default;
74
75
    static constexpr auto name = "array_map";
76
77
27
    static LambdaFunctionPtr create() { return std::make_shared<ArrayMapFunction>(); }
78
79
0
    std::string get_name() const override { return name; }
80
81
27
    Status prepare(RuntimeState* state, const VExprSPtrs& children) override {
82
27
        RETURN_IF_ERROR(LambdaFunction::prepare(state, children));
83
27
        DCHECK_GE(children.size(), 2);
84
85
27
        return _prepare_lambda_argument_binding(children[0], children.size() - 1,
86
27
                                                _lambda_argument_binding);
87
27
    }
88
89
    Status execute(VExprContext* context, const Block* block, const Selector* expr_selector,
90
                   size_t count, ColumnPtr& result_column, const DataTypePtr& result_type,
91
16
                   const VExprSPtrs& children) const override {
92
16
        LambdaArgs args_info;
93
94
        ///* array_map(lambda,arg1,arg2,.....) *///
95
        //1. child[1:end]->execute(src_block)
96
16
        ColumnsWithTypeAndName arguments(children.size() - 1);
97
34
        for (int i = 1; i < children.size(); ++i) {
98
18
            ColumnPtr column;
99
18
            RETURN_IF_ERROR(
100
18
                    children[i]->execute_column(context, block, expr_selector, count, column));
101
18
            arguments[i - 1].column = column;
102
18
            arguments[i - 1].type = children[i]->execute_type(block);
103
18
            arguments[i - 1].name = children[i]->expr_name();
104
18
        }
105
106
        // used for save column array outside null map
107
16
        auto outside_null_map = ColumnUInt8::create(
108
16
                arguments[0].column->convert_to_full_column_if_const()->size(), 0);
109
        // offset column
110
16
        MutableColumnPtr array_column_offset;
111
16
        size_t nested_array_column_rows = 0;
112
16
        ColumnPtr first_array_offsets = nullptr;
113
        //2. get the result column from executed expr, and the needed is nested column of array
114
16
        std::vector<ColumnPtr> lambda_datas(arguments.size());
115
16
        DataTypes lambda_argument_types(arguments.size());
116
117
34
        for (int i = 0; i < arguments.size(); ++i) {
118
18
            const auto& array_column_type_name = arguments[i];
119
18
            auto column_array = array_column_type_name.column->convert_to_full_column_if_const();
120
18
            auto type_array = array_column_type_name.type;
121
18
            if (type_array->is_nullable()) {
122
                // get the nullmap of nullable column
123
                // hold the null column instead of a reference 'cause `column_array` will be assigned and freed below.
124
1
                DORIS_CHECK(is_column_nullable(*column_array));
125
1
                auto column_array_nullmap =
126
1
                        assert_cast<const ColumnNullable&>(*column_array).get_null_map_column_ptr();
127
128
                // get the array column from nullable column
129
1
                column_array = assert_cast<const ColumnNullable*>(column_array.get())
130
1
                                       ->get_nested_column_ptr();
131
132
                // get the nested type from nullable type
133
1
                type_array = assert_cast<const DataTypeNullable*>(array_column_type_name.type.get())
134
1
                                     ->get_nested_type();
135
136
                // need to union nullmap from all columns
137
1
                VectorizedUtils::update_null_map(outside_null_map->get_data(),
138
1
                                                 column_array_nullmap->get_data());
139
1
            }
140
141
            // here is the array column
142
18
            const auto& col_array = assert_cast<const ColumnArray&>(*column_array);
143
144
18
            if (i == 0) {
145
16
                nested_array_column_rows = col_array.get_data_ptr()->size();
146
16
                first_array_offsets = col_array.get_offsets_ptr();
147
16
                const auto& off_data = col_array.get_offsets_column();
148
16
                array_column_offset = off_data.clone_resized(col_array.get_offsets_column().size());
149
16
                args_info.offsets_ptr = &col_array.get_offsets();
150
16
            } else {
151
                // select array_map((x,y)->x+y,c_array1,[0,1,2,3]) from array_test2;
152
                // c_array1: [0,1,2,3,4,5,6,7,8,9]
153
2
                const auto& array_offsets =
154
2
                        assert_cast<const ColumnArray::ColumnOffsets&>(*first_array_offsets)
155
2
                                .get_data();
156
2
                if (nested_array_column_rows != col_array.get_data_ptr()->size() ||
157
2
                    (!array_offsets.empty() &&
158
2
                     memcmp(array_offsets.data(), col_array.get_offsets().data(),
159
2
                            sizeof(array_offsets[0]) * array_offsets.size()) != 0)) {
160
0
                    return Status::InvalidArgument(
161
0
                            "in array map function, the input column size "
162
0
                            "are "
163
0
                            "not equal completely, nested column data rows 1st size is {}, {}th "
164
0
                            "size is {}.",
165
0
                            nested_array_column_rows, i + 1, col_array.get_data_ptr()->size());
166
0
                }
167
2
            }
168
18
            lambda_datas[i] = col_array.get_data_ptr();
169
18
            const auto& col_type = assert_cast<const DataTypeArray&>(*type_array);
170
18
            lambda_argument_types[i] = col_type.get_nested_type();
171
18
        }
172
16
        std::set<int> required_input_column_ids;
173
16
        children[0]->collect_slot_column_ids(required_input_column_ids);
174
16
        context->lambda_execution_context().collect_visible_binding_column_positions(
175
16
                required_input_column_ids);
176
16
        const int lambda_argument_base =
177
16
                required_input_column_ids.empty() ? 0 : *required_input_column_ids.rbegin() + 1;
178
16
        if (!_lambda_argument_binding.bind_by_name) {
179
1
            RETURN_IF_ERROR(
180
1
                    _set_legacy_lambda_argument_gap(children[0]->get_child(0), lambda_argument_base,
181
1
                                                    _lambda_argument_binding.argument_size));
182
1
        }
183
16
        std::vector<std::string> names(lambda_argument_base);
184
16
        DataTypes data_types(lambda_argument_base);
185
16
        std::vector<bool> materialized_input_columns(lambda_argument_base, false);
186
16
        names.reserve(lambda_argument_base + arguments.size());
187
16
        data_types.reserve(lambda_argument_base + arguments.size());
188
16
        for (int column_id : required_input_column_ids) {
189
12
            if (column_id < 0 || static_cast<size_t>(column_id) >= block->columns()) {
190
0
                return Status::InternalError(
191
0
                        "array_map lambda input column id {} is outside input block, block={}",
192
0
                        column_id, block->dump_structure());
193
0
            }
194
12
            materialized_input_columns[column_id] = true;
195
12
            names[column_id] = block->get_by_position(column_id).name;
196
12
            data_types[column_id] = block->get_by_position(column_id).type;
197
12
        }
198
37
        for (int i = 0; i < lambda_argument_base; ++i) {
199
21
            if (!materialized_input_columns[i]) {
200
                // Keep sparse input positions stable for SlotRef/parent lambda bindings without
201
                // materializing unrelated wide-table columns into every lambda batch.
202
9
                names[i] = "temp";
203
9
                data_types[i] = std::make_shared<DataTypeUInt8>();
204
9
            }
205
21
        }
206
34
        for (int i = 0; i < arguments.size(); ++i) {
207
18
            const auto& array_column_type_name = arguments[i];
208
18
            if (_lambda_argument_binding.bind_by_name &&
209
18
                i < _lambda_argument_binding.names.size()) {
210
16
                names.push_back(_lambda_argument_binding.names[i]);
211
16
            } else {
212
2
                names.push_back("R" + array_column_type_name.name);
213
2
            }
214
18
            data_types.push_back(lambda_argument_types[i]);
215
18
        }
216
217
16
        LambdaExecutionContext::Frame lambda_frame;
218
16
        lambda_frame.bind_by_name = _lambda_argument_binding.bind_by_name;
219
16
        lambda_frame.parent_bindings_visible = true;
220
33
        for (int i = 0; i < _lambda_argument_binding.argument_size; ++i) {
221
17
            const int column_position = lambda_argument_base + i;
222
17
            if (_lambda_argument_binding.bind_by_name) {
223
16
                lambda_frame.argument_bindings.push_back(
224
16
                        {_lambda_argument_binding.names[i], column_position});
225
16
            }
226
17
        }
227
16
        LambdaExecutionContext::FrameGuard lambda_frame_guard(context->lambda_execution_context(),
228
16
                                                              std::move(lambda_frame));
229
230
        // if column_array is NULL, we know the array_data_column will not write any data,
231
        // so the column is empty. eg : (x) -> concat('|',x + "1"). if still execute the lambda function, will cause the bolck rows are not equal
232
        // the x column is empty, but "|" is const literal, size of column is 1, so the block rows is 1, but the x column is empty, will be coredump.
233
18
        if (std::ranges::any_of(lambda_datas, [](const auto& v) { return v->empty(); })) {
234
0
            DataTypePtr nested_type;
235
0
            bool is_nullable = result_type->is_nullable();
236
0
            if (is_nullable) {
237
0
                nested_type =
238
0
                        assert_cast<const DataTypeNullable*>(result_type.get())->get_nested_type();
239
0
            } else {
240
0
                nested_type = result_type;
241
0
            }
242
0
            auto empty_nested_column = assert_cast<const DataTypeArray*>(nested_type.get())
243
0
                                               ->get_nested_type()
244
0
                                               ->create_column();
245
0
            auto result_array_column = ColumnArray::create(std::move(empty_nested_column),
246
0
                                                           std::move(array_column_offset));
247
248
0
            if (is_nullable) {
249
0
                result_column = ColumnNullable::create(std::move(result_array_column),
250
0
                                                       std::move(outside_null_map));
251
0
            } else {
252
0
                result_column = std::move(result_array_column);
253
0
            }
254
255
0
            return Status::OK();
256
0
        }
257
258
16
        MutableColumnPtr result_col = nullptr;
259
16
        DataTypePtr res_type;
260
261
        //process first row
262
16
        args_info.array_start = (*args_info.offsets_ptr)[args_info.current_row_idx - 1];
263
16
        args_info.cur_size =
264
16
                (*args_info.offsets_ptr)[args_info.current_row_idx] - args_info.array_start;
265
266
        // lambda block to exectute the lambda, and reuse the memory
267
16
        Block lambda_block;
268
16
        auto column_size = names.size();
269
16
        MutableColumns columns(column_size);
270
16
        do {
271
16
            bool mem_reuse = lambda_block.mem_reuse();
272
55
            for (int i = 0; i < column_size; i++) {
273
39
                if (mem_reuse) {
274
0
                    columns[i] = lambda_block.get_by_position(i).column->assert_mutable();
275
39
                } else {
276
39
                    columns[i] = data_types[i]->create_column();
277
39
                }
278
39
            }
279
            // batch_size of array nested data every time inorder to avoid memory overflow
280
30
            while (columns[lambda_argument_base]->size() < batch_size) {
281
30
                long max_step = batch_size - columns[lambda_argument_base]->size();
282
30
                long current_step = std::min(
283
30
                        max_step, (long)(args_info.cur_size - args_info.current_offset_in_array));
284
30
                size_t pos = args_info.array_start + args_info.current_offset_in_array;
285
64
                for (int i = 0; i < arguments.size() && current_step > 0; ++i) {
286
34
                    columns[lambda_argument_base + i]->insert_range_from(*lambda_datas[i], pos,
287
34
                                                                         current_step);
288
34
                }
289
30
                args_info.current_offset_in_array += current_step;
290
30
                args_info.current_repeat_times += current_step;
291
30
                if (args_info.current_offset_in_array >= args_info.cur_size) {
292
30
                    args_info.current_row_eos = true;
293
30
                }
294
30
                _repeat_input_columns(columns, block, args_info.current_repeat_times,
295
30
                                      materialized_input_columns, args_info.current_row_idx);
296
30
                args_info.current_repeat_times = 0;
297
30
                if (args_info.current_row_eos) {
298
                    //current row is end of array, move to next row
299
30
                    args_info.current_row_idx++;
300
30
                    args_info.current_offset_in_array = 0;
301
30
                    if (args_info.current_row_idx >= count) {
302
16
                        break;
303
16
                    }
304
14
                    args_info.current_row_eos = false;
305
14
                    args_info.array_start = (*args_info.offsets_ptr)[args_info.current_row_idx - 1];
306
14
                    args_info.cur_size = (*args_info.offsets_ptr)[args_info.current_row_idx] -
307
14
                                         args_info.array_start;
308
14
                }
309
30
            }
310
311
16
            if (!mem_reuse) {
312
55
                for (int i = 0; i < column_size; ++i) {
313
39
                    lambda_block.insert(
314
39
                            ColumnWithTypeAndName(std::move(columns[i]), data_types[i], names[i]));
315
39
                }
316
16
            }
317
            //3. child[0]->execute(new_block)
318
319
16
            ColumnPtr res_col;
320
            // lambda body executes on the internal lambda_block, not the original block.
321
            // The outer expr_selector is irrelevant here, so pass nullptr.
322
16
            RETURN_IF_ERROR(children[0]->execute_column(context, &lambda_block, nullptr,
323
16
                                                        lambda_block.rows(), res_col));
324
16
            res_col = res_col->convert_to_full_column_if_const();
325
16
            res_type = children[0]->execute_type(&lambda_block);
326
327
16
            if (!result_col) {
328
16
                result_col = res_col->clone_empty();
329
16
            }
330
16
            result_col->insert_range_from(*res_col, 0, res_col->size());
331
16
            lambda_block.clear_column_data(column_size);
332
16
        } while (args_info.current_row_idx < count);
333
334
        //4. get the result column after execution, reassemble it into a new array column, and return.
335
16
        if (result_type->is_nullable()) {
336
0
            if (res_type->is_nullable()) {
337
0
                result_column = ColumnNullable::create(
338
0
                        ColumnArray::create(std::move(result_col), std::move(array_column_offset)),
339
0
                        std::move(outside_null_map));
340
0
            } else {
341
                // deal with eg: select array_map(x -> x is null, [null, 1, 2]);
342
                // need to create the nested column null map for column array
343
0
                auto nested_null_map = ColumnUInt8::create(result_col->size(), 0);
344
345
0
                result_column = ColumnNullable::create(
346
0
                        ColumnArray::create(ColumnNullable::create(std::move(result_col),
347
0
                                                                   std::move(nested_null_map)),
348
0
                                            std::move(array_column_offset)),
349
0
                        std::move(outside_null_map));
350
0
            }
351
16
        } else {
352
16
            if (res_type->is_nullable()) {
353
4
                result_column =
354
4
                        ColumnArray::create(std::move(result_col), std::move(array_column_offset));
355
12
            } else {
356
12
                auto nested_null_map = ColumnUInt8::create(result_col->size(), 0);
357
358
12
                result_column = ColumnArray::create(
359
12
                        ColumnNullable::create(std::move(result_col), std::move(nested_null_map)),
360
12
                        std::move(array_column_offset));
361
12
            }
362
16
        }
363
16
        return Status::OK();
364
16
    }
365
366
private:
367
    struct LambdaArgumentBinding {
368
        bool bind_by_name = true;
369
        size_t argument_size = 0;
370
        std::vector<std::string> names;
371
    };
372
373
    Status _prepare_lambda_argument_binding(const VExprSPtr& expr, size_t expected_argument_size,
374
27
                                            LambdaArgumentBinding& argument_binding) const {
375
27
        DORIS_CHECK_EQ(expr->node_type(), TExprNodeType::LAMBDA_FUNCTION_EXPR);
376
27
        const auto* lambda_expr = assert_cast<const VLambdaFunctionExpr*>(expr.get());
377
378
27
        argument_binding.argument_size = 0;
379
27
        argument_binding.names.clear();
380
27
        argument_binding.bind_by_name = lambda_expr->has_argument_names();
381
382
27
        if (!argument_binding.bind_by_name) {
383
5
            if (_contains_nested_lambda_call(expr->get_child(0))) {
384
1
                return Status::InternalError(
385
1
                        "Cannot resolve nested lambda argument without lambda metadata");
386
1
            }
387
4
            argument_binding.argument_size = expected_argument_size;
388
4
            argument_binding.names.resize(expected_argument_size);
389
4
            return Status::OK();
390
5
        }
391
392
22
        argument_binding.names = lambda_expr->argument_names();
393
22
        if (argument_binding.names.size() > expected_argument_size) {
394
0
            return Status::InternalError(
395
0
                    "lambda argument metadata size exceeds parameter size, maximum={}, actual={}",
396
0
                    expected_argument_size, argument_binding.names.size());
397
0
        }
398
22
        argument_binding.argument_size = argument_binding.names.size();
399
22
        if (std::ranges::any_of(argument_binding.names,
400
23
                                [](const auto& argument_name) { return argument_name.empty(); })) {
401
0
            return Status::InternalError("lambda argument metadata contains empty name");
402
0
        }
403
22
        return Status::OK();
404
22
    }
405
406
    Status _set_legacy_lambda_argument_gap(const VExprSPtr& expr, int lambda_argument_base,
407
1
                                           size_t argument_size) const {
408
1
        if (expr->is_column_ref()) {
409
1
            auto* ref = static_cast<VColumnRef*>(expr.get());
410
1
            DORIS_CHECK_GE(ref->column_id(), 0);
411
1
            DORIS_CHECK_LT(static_cast<size_t>(ref->column_id()), argument_size);
412
1
            const int argument_index = ref->column_id();
413
1
            ref->set_gap(lambda_argument_base + argument_index - ref->column_id());
414
1
        } else {
415
0
            for (const auto& child : expr->children()) {
416
0
                RETURN_IF_ERROR(_set_legacy_lambda_argument_gap(child, lambda_argument_base,
417
0
                                                                argument_size));
418
0
            }
419
0
        }
420
1
        return Status::OK();
421
1
    }
422
423
9
    bool _is_lambda_call_with_lambda_expr(const VExprSPtr& expr) const {
424
9
        return expr->node_type() == TExprNodeType::LAMBDA_FUNCTION_CALL_EXPR &&
425
9
               !expr->children().empty() &&
426
9
               expr->children()[0]->node_type() == TExprNodeType::LAMBDA_FUNCTION_EXPR;
427
9
    }
428
429
9
    bool _contains_nested_lambda_call(const VExprSPtr& expr) const {
430
9
        if (_is_lambda_call_with_lambda_expr(expr)) {
431
1
            return true;
432
1
        }
433
8
        return std::ranges::any_of(expr->children(), [this](const auto& child) {
434
4
            return _contains_nested_lambda_call(child);
435
4
        });
436
9
    }
437
438
    void _repeat_input_columns(std::vector<MutableColumnPtr>& columns, const Block* block,
439
                               int repeat_times,
440
                               const std::vector<bool>& materialized_input_columns,
441
30
                               int64_t row_idx) const {
442
30
        if (!repeat_times || materialized_input_columns.empty()) {
443
10
            return;
444
10
        }
445
71
        for (size_t i = 0; i < materialized_input_columns.size(); i++) {
446
51
            if (!materialized_input_columns[i]) {
447
23
                columns[i]->resize(columns[i]->size() + repeat_times);
448
23
                continue;
449
23
            }
450
28
            DORIS_CHECK(block != nullptr);
451
28
            auto src_column = block->get_by_position(i).column->convert_to_full_column_if_const();
452
28
            if (check_and_get_column<ColumnNothing>(src_column.get())) {
453
                // A ColumnNothing in the outer block is a placeholder for an unmaterialized
454
                // virtual column. Keep it as a placeholder in the lambda block as well, so
455
                // VirtualSlotRef can still materialize it lazily if the lambda body reads it.
456
0
                if (!check_and_get_column<ColumnNothing>(columns[i].get())) {
457
0
                    columns[i] = ColumnNothing::create(columns[i]->size());
458
0
                }
459
0
            }
460
28
            columns[i]->insert_many_from(*src_column, row_idx, repeat_times);
461
28
        }
462
20
    }
463
464
    LambdaArgumentBinding _lambda_argument_binding;
465
};
466
467
1
void register_function_array_map(doris::LambdaFunctionFactory& factory) {
468
1
    factory.register_function<ArrayMapFunction>();
469
1
}
470
471
} // namespace doris