This is an automated email from the ASF dual-hosted git repository.
HappenLee pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/doris.git
The following commit(s) were added to refs/heads/master by this push:
new cb0f7a434bd [Opt](lambda) Aviod redundant copies during lambda execute
(#66590)
cb0f7a434bd is described below
commit cb0f7a434bdae0b9818646d42a8491030d8ecbae
Author: linrrarity <[email protected]>
AuthorDate: Tue Aug 18 18:17:16 2026 +0800
[Opt](lambda) Aviod redundant copies during lambda execute (#66590)
`array_map` flattens array elements into an internal block before
executing its lambda expression. The previous implementation introduced
several redundant copies and materializations on this path:
- Nested array elements were copied into new columns even when all
elements could be processed in a single lambda batch.
- Captured outer columns were expanded repeatedly for each outer row.
- Constant captured columns were materialized as full columns.
- Sparse, unused input positions could create unnecessary placeholder
column work.
- The first lambda result batch was copied into a newly created result
column.
- Internal lambda execution could receive the outer selector even though
its input block had already been selected and flattened.
These operations increased CPU usage and peak memory consumption,
especially for large arrays, wide captured inputs, and complex lambda
expressions.
This change optimizes the `array_map` execution path by:
- Reusing the nested array data columns directly when all nested rows
fit within one lambda batch.
- Preserving captured constant columns as `ColumnConst`.
- Expanding row-dependent captured columns in one operation with
selector-aware source row indices.
- Representing unused sparse input positions with `ColumnNothing`.
- Reusing lambda block column allocations across batches.
- Taking ownership of the first lambda result column instead of cloning
and copying it.
- Passing no outer selector when executing the already-materialized
internal lambda block.
- Centralizing nullable array result construction to keep the execution
path simpler.
- Sizing internal lambda batches according to the runtime row and byte
budgets.
The optimization preserves array element order, nullable semantics,
captured-column mapping, nested lambda behavior, and SQL result
compatibility.
Performance:
```text
Before:
Doris> select sum(array_sum(array_map(x -> x + 1, array(id % 5, id % 7, id
% 11)))) from t_explode;
+-----------------------------------------------------------------------+
| sum(array_sum(array_map(x -> x + 1, array(id % 5, id % 7, id % 11)))) |
+-----------------------------------------------------------------------+
| 6499979 |
+-----------------------------------------------------------------------+
1 row in set (0.082 sec)
After:
Doris> select sum(array_sum(array_map(x -> x + 1, array(id % 5, id % 7, id
% 11)))) from t_explode;
+-----------------------------------------------------------------------+
| sum(array_sum(array_map(x -> x + 1, array(id % 5, id % 7, id % 11)))) |
+-----------------------------------------------------------------------+
| 6499979 |
+-----------------------------------------------------------------------+
1 row in set (0.056 sec)
Before:
Doris> select sum(array_sum(array_map(x -> x + id, array(id % 5, id % 7, id
% 11)))) from t_explode;
+------------------------------------------------------------------------+
| sum(array_sum(array_map(x -> x + id, array(id % 5, id % 7, id % 11)))) |
+------------------------------------------------------------------------+
| 375004249979 |
+------------------------------------------------------------------------+
1 row in set (0.091 sec)
After:
Doris> select sum(array_sum(array_map(x -> x + id, array(id % 5, id % 7, id
% 11)))) from t_explode;
+------------------------------------------------------------------------+
| sum(array_sum(array_map(x -> x + id, array(id % 5, id % 7, id % 11)))) |
+------------------------------------------------------------------------+
| 375004249979 |
+------------------------------------------------------------------------+
1 row in set (0.070 sec)
Before:
Doris> select sum(array_sum(array_map(x -> x + 1, arr))) from t_explode;
+--------------------------------------------+
| sum(array_sum(array_map(x -> x + 1, arr))) |
+--------------------------------------------+
| 380997500000 |
+--------------------------------------------+
1 row in set (6.759 sec)
After:
Doris> select sum(array_sum(array_map(x -> x + 1, arr))) from t_explode;
+--------------------------------------------+
| sum(array_sum(array_map(x -> x + 1, arr))) |
+--------------------------------------------+
| 380997500000 |
+--------------------------------------------+
1 row in set (5.807 sec)
```
---
be/src/exprs/lambda_function/lambda_function.h | 5 +-
.../exprs/lambda_function/varray_map_function.cpp | 251 ++++++--
.../lambda_function/array_map_function_test.cpp | 679 ++++++++++++++++++++-
.../test_array_map_function_with_column.out | 14 +
.../test_array_map_function_with_column.groovy | 25 +
5 files changed, 921 insertions(+), 53 deletions(-)
diff --git a/be/src/exprs/lambda_function/lambda_function.h
b/be/src/exprs/lambda_function/lambda_function.h
index 31976617a54..8fbfa6ac996 100644
--- a/be/src/exprs/lambda_function/lambda_function.h
+++ b/be/src/exprs/lambda_function/lambda_function.h
@@ -31,8 +31,7 @@ public:
virtual std::string get_name() const = 0;
- virtual doris::Status prepare(RuntimeState* state, const VExprSPtrs&
children) {
- batch_size = state->batch_size();
+ virtual doris::Status prepare(RuntimeState* /*state*/, const VExprSPtrs&
/*children*/) {
return Status::OK();
}
@@ -40,8 +39,6 @@ public:
const Selector* selector, size_t count,
ColumnPtr& result_column,
const DataTypePtr& result_type,
const VExprSPtrs& children) const = 0;
-
- int batch_size;
};
using LambdaFunctionPtr = std::shared_ptr<LambdaFunction>;
diff --git a/be/src/exprs/lambda_function/varray_map_function.cpp
b/be/src/exprs/lambda_function/varray_map_function.cpp
index 1c22c23b7c8..da0722ee0ad 100644
--- a/be/src/exprs/lambda_function/varray_map_function.cpp
+++ b/be/src/exprs/lambda_function/varray_map_function.cpp
@@ -16,6 +16,7 @@
// under the License.
#include <algorithm>
+#include <limits>
#include <memory>
#include <set>
#include <string>
@@ -30,6 +31,7 @@
#include "core/block/columns_with_type_and_name.h"
#include "core/column/column.h"
#include "core/column/column_array.h"
+#include "core/column/column_const.h"
#include "core/column/column_nothing.h"
#include "core/column/column_nullable.h"
#include "core/column/column_vector.h"
@@ -45,6 +47,7 @@
#include "exprs/vcolumn_ref.h"
#include "exprs/vexpr_context.h"
#include "exprs/vlambda_function_expr.h"
+#include "util/block_budget.h"
namespace doris {
@@ -60,8 +63,6 @@ struct LambdaArgs {
int64_t cur_size = 0;
// offset of column array
const ColumnArray::Offsets64* offsets_ptr = nullptr;
- // expend data of repeat times
- int current_repeat_times = 0;
// whether the current row of the original block has been extended
bool current_row_eos = false;
};
@@ -82,6 +83,9 @@ public:
RETURN_IF_ERROR(LambdaFunction::prepare(state, children));
DCHECK_GE(children.size(), 2);
+ _lambda_block_budget =
+ BlockBudget(state->batch_size(),
state->preferred_block_size_bytes());
+
return _prepare_lambda_argument_binding(children[0], children.size() -
1,
_lambda_argument_binding);
}
@@ -183,6 +187,7 @@ public:
std::vector<std::string> names(lambda_argument_base);
DataTypes data_types(lambda_argument_base);
std::vector<bool> materialized_input_columns(lambda_argument_base,
false);
+ bool has_row_dependent_captures = false;
names.reserve(lambda_argument_base + arguments.size());
data_types.reserve(lambda_argument_base + arguments.size());
for (int column_id : required_input_column_ids) {
@@ -194,6 +199,9 @@ public:
materialized_input_columns[column_id] = true;
names[column_id] = block->get_by_position(column_id).name;
data_types[column_id] = block->get_by_position(column_id).type;
+ const auto& input_column =
block->get_by_position(column_id).column;
+ has_row_dependent_captures |= !is_column_const(*input_column) &&
+
!check_and_get_column<ColumnNothing>(input_column.get());
}
for (int i = 0; i < lambda_argument_base; ++i) {
if (!materialized_input_columns[i]) {
@@ -251,7 +259,71 @@ public:
} else {
result_column = std::move(result_array_column);
}
+ return Status::OK();
+ }
+
+ const size_t lambda_batch_rows =
+ _calculate_lambda_batch_size(children[0], lambda_datas, block,
+ required_input_column_ids,
has_row_dependent_captures);
+
+ // Lambda arguments are already stored contiguously in the input
arrays. When all nested
+ // rows fit in one lambda batch, reuse those columns directly and only
materialize captured
+ // outer columns whose values depend on the outer row.
+ if (nested_array_column_rows > 0 && nested_array_column_rows <=
lambda_batch_rows) {
+ Block lambda_block;
+ PaddedPODArray<IColumn::ColumnIndex> captured_source_row_indices;
+ MutableColumns captured_columns(lambda_argument_base);
+ for (int i = 0; i < lambda_argument_base; ++i) {
+ if (!materialized_input_columns[i]) {
+ captured_columns[i] =
ColumnNothing::create(nested_array_column_rows);
+ continue;
+ }
+
+ const auto& source_column = block->get_by_position(i).column;
+ if (is_column_const(*source_column)) {
+ captured_columns[i] =
source_column->clone_resized(nested_array_column_rows);
+ } else if
(check_and_get_column<ColumnNothing>(source_column.get())) {
+ captured_columns[i] =
ColumnNothing::create(nested_array_column_rows);
+ } else {
+ if (captured_source_row_indices.empty()) {
+
captured_source_row_indices.reserve(nested_array_column_rows);
+ size_t previous_offset = 0;
+ for (size_t row_idx = 0; row_idx < count; ++row_idx) {
+ const size_t current_offset =
(*args_info.offsets_ptr)[row_idx];
+ const size_t repeat_times = current_offset -
previous_offset;
+ const auto source_row =
+ expr_selector == nullptr
+ ?
static_cast<IColumn::ColumnIndex>(row_idx)
+ : (*expr_selector)[row_idx];
+
_append_captured_source_row_indices(captured_source_row_indices,
+ source_row,
repeat_times);
+ previous_offset = current_offset;
+ }
+ }
+ captured_columns[i] = data_types[i]->create_column();
+ captured_columns[i]->insert_indices_from(
+ *source_column, captured_source_row_indices.data(),
+ captured_source_row_indices.data() +
+ captured_source_row_indices.size());
+ }
+ }
+ for (int i = 0; i < lambda_argument_base; ++i) {
+
lambda_block.insert(ColumnWithTypeAndName(std::move(captured_columns[i]),
+ data_types[i],
names[i]));
+ }
+ for (int i = 0; i < arguments.size(); ++i) {
+ lambda_block.insert(ColumnWithTypeAndName(lambda_datas[i],
lambda_argument_types[i],
+
names[lambda_argument_base + i]));
+ }
+ ColumnPtr res_col;
+ RETURN_IF_ERROR(children[0]->execute_column(context,
&lambda_block, nullptr,
+
nested_array_column_rows, res_col));
+ res_col = res_col->convert_to_full_column_if_const();
+ auto res_type = children[0]->execute_type(&lambda_block);
+ result_column =
+ _create_result_column(std::move(res_col),
std::move(array_column_offset),
+ std::move(outside_null_map),
res_type, result_type);
return Status::OK();
}
@@ -267,18 +339,28 @@ public:
Block lambda_block;
auto column_size = names.size();
MutableColumns columns(column_size);
+ PaddedPODArray<IColumn::ColumnIndex> captured_source_row_indices;
+ if (has_row_dependent_captures) {
+ captured_source_row_indices.reserve(lambda_batch_rows);
+ }
do {
+ captured_source_row_indices.clear();
bool mem_reuse = lambda_block.mem_reuse();
for (int i = 0; i < column_size; i++) {
if (mem_reuse) {
columns[i] =
lambda_block.get_by_position(i).column->assert_mutable();
+ } else if (i < lambda_argument_base &&
!materialized_input_columns[i]) {
+ columns[i] = ColumnNothing::create(0);
+ } else if (i < lambda_argument_base &&
materialized_input_columns[i] &&
+ is_column_const(*block->get_by_position(i).column))
{
+ columns[i] =
block->get_by_position(i).column->clone_resized(0);
} else {
columns[i] = data_types[i]->create_column();
}
}
- // batch_size of array nested data every time inorder to avoid
memory overflow
- while (columns[lambda_argument_base]->size() < batch_size) {
- long max_step = batch_size -
columns[lambda_argument_base]->size();
+ // lambda_batch_rows of array nested data every time inorder to
avoid memory overflow
+ while (columns[lambda_argument_base]->size() < lambda_batch_rows) {
+ long max_step = lambda_batch_rows -
columns[lambda_argument_base]->size();
long current_step = std::min(
max_step, (long)(args_info.cur_size -
args_info.current_offset_in_array));
size_t pos = args_info.array_start +
args_info.current_offset_in_array;
@@ -287,13 +369,17 @@ public:
current_step);
}
args_info.current_offset_in_array += current_step;
- args_info.current_repeat_times += current_step;
+ if (has_row_dependent_captures) {
+ const auto source_row =
+ expr_selector == nullptr
+ ?
static_cast<IColumn::ColumnIndex>(args_info.current_row_idx)
+ :
(*expr_selector)[args_info.current_row_idx];
+
_append_captured_source_row_indices(captured_source_row_indices, source_row,
+ current_step);
+ }
if (args_info.current_offset_in_array >= args_info.cur_size) {
args_info.current_row_eos = true;
}
- _repeat_input_columns(columns, block,
args_info.current_repeat_times,
- materialized_input_columns,
args_info.current_row_idx);
- args_info.current_repeat_times = 0;
if (args_info.current_row_eos) {
//current row is end of array, move to next row
args_info.current_row_idx++;
@@ -307,6 +393,9 @@ public:
args_info.array_start;
}
}
+ const size_t current_lambda_batch_rows =
columns[lambda_argument_base]->size();
+ _repeat_input_columns(columns, block, captured_source_row_indices,
+ materialized_input_columns,
current_lambda_batch_rows);
if (!mem_reuse) {
for (int i = 0; i < column_size; ++i) {
@@ -325,45 +414,102 @@ public:
res_type = children[0]->execute_type(&lambda_block);
if (!result_col) {
- result_col = res_col->clone_empty();
+ result_col = IColumn::mutate(std::move(res_col));
+ } else {
+ result_col->insert_range_from(*res_col, 0, res_col->size());
}
- result_col->insert_range_from(*res_col, 0, res_col->size());
lambda_block.clear_column_data(column_size);
} while (args_info.current_row_idx < count);
//4. get the result column after execution, reassemble it into a new
array column, and return.
- if (result_type->is_nullable()) {
- if (res_type->is_nullable()) {
- result_column = ColumnNullable::create(
- ColumnArray::create(std::move(result_col),
std::move(array_column_offset)),
- std::move(outside_null_map));
- } else {
- // deal with eg: select array_map(x -> x is null, [null, 1,
2]);
- // need to create the nested column null map for column array
- auto nested_null_map = ColumnUInt8::create(result_col->size(),
0);
-
- result_column = ColumnNullable::create(
-
ColumnArray::create(ColumnNullable::create(std::move(result_col),
-
std::move(nested_null_map)),
- std::move(array_column_offset)),
- std::move(outside_null_map));
- }
- } else {
- if (res_type->is_nullable()) {
- result_column =
- ColumnArray::create(std::move(result_col),
std::move(array_column_offset));
- } else {
- auto nested_null_map = ColumnUInt8::create(result_col->size(),
0);
+ result_column = _create_result_column(std::move(result_col),
std::move(array_column_offset),
+ std::move(outside_null_map),
res_type, result_type);
+ return Status::OK();
+ }
+
+private:
+ static bool _has_variable_length_column(const VExprSPtr& expr) {
+ return !expr->data_type()->have_maximum_size_of_value() ||
+ std::ranges::any_of(expr->children(), [](const auto& child) {
+ return _has_variable_length_column(child);
+ });
+ }
- result_column = ColumnArray::create(
- ColumnNullable::create(std::move(result_col),
std::move(nested_null_map)),
- std::move(array_column_offset));
+ // A referenced non-const capture is expanded once for every nested array
element before
+ // lambda evaluation. Expanding the full nested cardinality at once can
create multi-gigabyte
+ // temporary columns (for example, a 5,000-byte VARCHAR repeated 1,000,000
times) and exceed
+ // ColumnString's UInt32 offset limit. Keep capture expansion and lambda
evaluation within the
+ // runtime block budget, while retaining direct nested-input reuse when
one batch is sufficient.
+ // Rule: use the external row budget if any lambda input, output, or
intermediate column
+ // is variable-length. Fixed-width columns have predictable memory usage,
so also apply
+ // the external byte budget to their estimated bytes per row.
+ size_t _calculate_lambda_batch_size(const VExprSPtr& lambda_expr,
+ const std::vector<ColumnPtr>&
lambda_datas,
+ const Block* block,
+ const std::set<int>&
required_input_column_ids,
+ bool has_row_dependent_captures) const
{
+ const auto add_bytes_with_saturation = [](size_t current_bytes, size_t
additional_bytes) {
+ constexpr size_t max_bytes = std::numeric_limits<size_t>::max();
+ return additional_bytes > max_bytes - current_bytes ? max_bytes
+ :
current_bytes + additional_bytes;
+ };
+
+ if (_has_variable_length_column(lambda_expr)) {
+ return _lambda_block_budget.max_rows;
+ }
+
+ size_t estimated_lambda_bytes_per_row =
lambda_expr->estimate_memory(1);
+ for (const auto& lambda_data : lambda_datas) {
+ estimated_lambda_bytes_per_row = add_bytes_with_saturation(
+ estimated_lambda_bytes_per_row,
lambda_data->get_max_row_byte_size());
+ }
+ if (has_row_dependent_captures) {
+ estimated_lambda_bytes_per_row = add_bytes_with_saturation(
+ estimated_lambda_bytes_per_row,
sizeof(IColumn::ColumnIndex));
+ for (int column_id : required_input_column_ids) {
+ const auto& input_column =
block->get_by_position(column_id).column;
+ if (!is_column_const(*input_column) &&
+ !check_and_get_column<ColumnNothing>(input_column.get())) {
+ estimated_lambda_bytes_per_row = add_bytes_with_saturation(
+ estimated_lambda_bytes_per_row,
input_column->get_max_row_byte_size());
+ }
}
}
- return Status::OK();
+ return
_lambda_block_budget.effective_max_rows(estimated_lambda_bytes_per_row);
+ }
+
+ static void _append_captured_source_row_indices(
+ PaddedPODArray<IColumn::ColumnIndex>& captured_source_row_indices,
+ IColumn::ColumnIndex source_row, size_t repeat_times) {
+ const size_t old_size = captured_source_row_indices.size();
+ captured_source_row_indices.resize(old_size + repeat_times);
+ std::fill(captured_source_row_indices.begin() + old_size,
captured_source_row_indices.end(),
+ source_row);
+ }
+
+ static ColumnPtr _create_result_column(ColumnPtr result_col,
+ MutableColumnPtr
array_column_offset,
+ MutableColumnPtr outside_null_map,
+ const DataTypePtr& res_type,
+ const DataTypePtr& result_type) {
+ ColumnPtr nested_column = std::move(result_col);
+ if (!res_type->is_nullable()) {
+ // deal with eg: select array_map(x -> x is null, [null, 1, 2]);
+ // need to create the nested column null map for column array
+ auto nested_null_map = ColumnUInt8::create(nested_column->size(),
0);
+ nested_column =
+ ColumnNullable::create(std::move(nested_column),
std::move(nested_null_map));
+ }
+
+ auto result_array_column =
+ ColumnArray::create(std::move(nested_column),
std::move(array_column_offset));
+ if (result_type->is_nullable()) {
+ return ColumnNullable::create(std::move(result_array_column),
+ std::move(outside_null_map));
+ }
+ return result_array_column;
}
-private:
struct LambdaArgumentBinding {
bool bind_by_name = true;
size_t argument_size = 0;
@@ -435,33 +581,42 @@ private:
});
}
- void _repeat_input_columns(std::vector<MutableColumnPtr>& columns, const
Block* block,
- int repeat_times,
- const std::vector<bool>&
materialized_input_columns,
- int64_t row_idx) const {
- if (!repeat_times || materialized_input_columns.empty()) {
+ void _repeat_input_columns(
+ std::vector<MutableColumnPtr>& columns, const Block* block,
+ const PaddedPODArray<IColumn::ColumnIndex>&
captured_source_row_indices,
+ const std::vector<bool>& materialized_input_columns, size_t
lambda_batch_rows) const {
+ if (lambda_batch_rows == 0 || materialized_input_columns.empty()) {
return;
}
for (size_t i = 0; i < materialized_input_columns.size(); i++) {
if (!materialized_input_columns[i]) {
- columns[i]->resize(columns[i]->size() + repeat_times);
+ columns[i]->resize(lambda_batch_rows);
continue;
}
DORIS_CHECK(block != nullptr);
- auto src_column =
block->get_by_position(i).column->convert_to_full_column_if_const();
- if (check_and_get_column<ColumnNothing>(src_column.get())) {
+ const auto& src_column = block->get_by_position(i).column;
+ if (is_column_const(*src_column)) {
+ columns[i]->resize(lambda_batch_rows);
+ } else if (check_and_get_column<ColumnNothing>(src_column.get())) {
// A ColumnNothing in the outer block is a placeholder for an
unmaterialized
// virtual column. Keep it as a placeholder in the lambda
block as well, so
// VirtualSlotRef can still materialize it lazily if the
lambda body reads it.
if (!check_and_get_column<ColumnNothing>(columns[i].get())) {
- columns[i] = ColumnNothing::create(columns[i]->size());
+ columns[i] = ColumnNothing::create(lambda_batch_rows);
+ } else {
+ columns[i]->resize(lambda_batch_rows);
}
+ } else {
+ DCHECK_EQ(captured_source_row_indices.size(),
lambda_batch_rows);
+ columns[i]->insert_indices_from(
+ *src_column, captured_source_row_indices.data(),
+ captured_source_row_indices.data() +
captured_source_row_indices.size());
}
- columns[i]->insert_many_from(*src_column, row_idx, repeat_times);
}
}
LambdaArgumentBinding _lambda_argument_binding;
+ BlockBudget _lambda_block_budget {1, 0};
};
void register_function_array_map(doris::LambdaFunctionFactory& factory) {
diff --git a/be/test/exprs/lambda_function/array_map_function_test.cpp
b/be/test/exprs/lambda_function/array_map_function_test.cpp
index 42469b52c3b..5e0e5da0781 100644
--- a/be/test/exprs/lambda_function/array_map_function_test.cpp
+++ b/be/test/exprs/lambda_function/array_map_function_test.cpp
@@ -19,19 +19,24 @@
#include <gen_cpp/Types_types.h>
#include <gtest/gtest.h>
+#include <algorithm>
#include <memory>
#include <string>
#include <vector>
+#include "common/config.h"
#include "core/assert_cast.h"
#include "core/block/block.h"
#include "core/column/column_array.h"
#include "core/column/column_const.h"
+#include "core/column/column_nothing.h"
#include "core/column/column_nullable.h"
+#include "core/column/column_string.h"
#include "core/column/column_vector.h"
#include "core/data_type/data_type_array.h"
#include "core/data_type/data_type_nullable.h"
#include "core/data_type/data_type_number.h"
+#include "core/data_type/data_type_string.h"
#include "exprs/vcolumn_ref.h"
#include "exprs/vexpr_context.h"
#include "exprs/vlambda_function_call_expr.h"
@@ -39,6 +44,7 @@
#include "exprs/vslot_ref.h"
#include "runtime/descriptors.h"
#include "runtime/runtime_state.h"
+#include "util/defer_op.h"
namespace doris {
@@ -112,6 +118,163 @@ private:
std::string _name;
};
+class MockCapturedInputExpr final : public VExpr {
+public:
+ MockCapturedInputExpr(DataTypePtr type, const IColumn* expected_input,
bool* used_direct_input,
+ bool* kept_capture_const)
+ : VExpr(type, false),
+ _type(std::move(type)),
+ _expected_input(expected_input),
+ _used_direct_input(used_direct_input),
+ _kept_capture_const(kept_capture_const) {}
+
+ const std::string& expr_name() const override { return _name; }
+
+ Status execute_column_impl(VExprContext* context, const Block* block,
const Selector* selector,
+ size_t count, ColumnPtr& result_column) const
override {
+ *_kept_capture_const =
is_column_const(*block->get_by_position(0).column);
+ *_used_direct_input = block->get_by_position(1).column.get() ==
_expected_input;
+
+ ColumnPtr captured;
+ ColumnPtr input;
+ RETURN_IF_ERROR(get_child(0)->execute_column(context, block, selector,
count, captured));
+ RETURN_IF_ERROR(get_child(1)->execute_column(context, block, selector,
count, input));
+ captured = captured->convert_to_full_column_if_const();
+ input = input->convert_to_full_column_if_const();
+
+ const IColumn* captured_data = captured.get();
+ if (const auto* nullable =
check_and_get_column<ColumnNullable>(captured_data)) {
+ captured_data = &nullable->get_nested_column();
+ }
+ const IColumn* input_data = input.get();
+ if (const auto* nullable =
check_and_get_column<ColumnNullable>(input_data)) {
+ input_data = &nullable->get_nested_column();
+ }
+ const auto& captured_values = assert_cast<const
ColumnInt32&>(*captured_data);
+ const auto& input_values = assert_cast<const
ColumnInt32&>(*input_data);
+ auto result = ColumnInt32::create();
+ result->reserve(count);
+ for (size_t i = 0; i < count; ++i) {
+ result->insert_value(captured_values.get_element(i) +
input_values.get_element(i));
+ }
+ result_column = std::move(result);
+ return Status::OK();
+ }
+
+ DataTypePtr execute_type(const Block* /*block*/) const override { return
_type; }
+
+private:
+ DataTypePtr _type;
+ const IColumn* _expected_input;
+ bool* _used_direct_input;
+ bool* _kept_capture_const;
+ std::string _name = "mock_captured_const_direct";
+};
+
+class MockIncrementExpr final : public VExpr {
+public:
+ explicit MockIncrementExpr(DataTypePtr type) : VExpr(type, false),
_type(std::move(type)) {}
+
+ const std::string& expr_name() const override { return _name; }
+
+ Status execute_column_impl(VExprContext* context, const Block* block,
const Selector* selector,
+ size_t count, ColumnPtr& result_column) const
override {
+ ColumnPtr input;
+ RETURN_IF_ERROR(get_child(0)->execute_column(context, block, selector,
count, input));
+ const IColumn* data = input.get();
+ if (const auto* nullable = check_and_get_column<ColumnNullable>(data))
{
+ data = &nullable->get_nested_column();
+ }
+ const auto& values = assert_cast<const ColumnInt32&>(*data);
+ auto result = ColumnInt32::create();
+ result->reserve(count);
+ for (size_t i = 0; i < count; ++i) {
+ result->insert_value(values.get_element(i) + 1);
+ }
+ result_column = std::move(result);
+ return Status::OK();
+ }
+
+ DataTypePtr execute_type(const Block* /*block*/) const override { return
_type; }
+
+private:
+ DataTypePtr _type;
+ std::string _name = "mock_increment";
+};
+
+class MockBatchSizeExpr final : public VExpr {
+public:
+ MockBatchSizeExpr(DataTypePtr type, std::vector<size_t>*
observed_batch_sizes)
+ : VExpr(type, false),
+ _type(std::move(type)),
+ _observed_batch_sizes(observed_batch_sizes) {}
+
+ const std::string& expr_name() const override { return _name; }
+
+ Status execute_column_impl(VExprContext* /*context*/, const Block*
/*block*/,
+ const Selector* /*selector*/, size_t count,
+ ColumnPtr& result_column) const override {
+ _observed_batch_sizes->push_back(count);
+ result_column = ColumnInt32::create(count, 0);
+ return Status::OK();
+ }
+
+ DataTypePtr execute_type(const Block* /*block*/) const override { return
_type; }
+
+private:
+ DataTypePtr _type;
+ std::vector<size_t>* _observed_batch_sizes;
+ std::string _name = "mock_batch_size";
+};
+
+class MockGreatestExpr final : public VExpr {
+public:
+ MockGreatestExpr(DataTypePtr type, std::vector<size_t>*
observed_batch_sizes)
+ : VExpr(type, false),
+ _type(std::move(type)),
+ _observed_batch_sizes(observed_batch_sizes) {}
+
+ const std::string& expr_name() const override { return _name; }
+
+ Status execute_column_impl(VExprContext* context, const Block* block,
const Selector* selector,
+ size_t count, ColumnPtr& result_column) const
override {
+ _observed_batch_sizes->push_back(count);
+ std::vector<ColumnPtr> inputs;
+ inputs.reserve(children().size());
+ for (const auto& child : children()) {
+ ColumnPtr input;
+ RETURN_IF_ERROR(child->execute_column(context, block, selector,
count, input));
+ inputs.push_back(input->convert_to_full_column_if_const());
+ }
+
+ auto result = ColumnInt32::create();
+ result->reserve(count);
+ for (size_t row = 0; row < count; ++row) {
+ int32_t greatest = _get_int_data(inputs[0]).get_element(row);
+ for (size_t i = 1; i < inputs.size(); ++i) {
+ greatest = std::max(greatest,
_get_int_data(inputs[i]).get_element(row));
+ }
+ result->insert_value(greatest);
+ }
+ result_column = std::move(result);
+ return Status::OK();
+ }
+
+ DataTypePtr execute_type(const Block* /*block*/) const override { return
_type; }
+
+private:
+ const ColumnInt32& _get_int_data(const ColumnPtr& column) const {
+ if (const auto* nullable =
check_and_get_column<ColumnNullable>(column.get())) {
+ return assert_cast<const
ColumnInt32&>(nullable->get_nested_column());
+ }
+ return assert_cast<const ColumnInt32&>(*column);
+ }
+
+ DataTypePtr _type;
+ std::vector<size_t>* _observed_batch_sizes;
+ std::string _name = "mock_greatest";
+};
+
class MockSubtractExpr final : public VExpr {
public:
explicit MockSubtractExpr(DataTypePtr type) : VExpr(type, false),
_type(std::move(type)) {}
@@ -153,12 +316,18 @@ private:
class MockAddExpr final : public VExpr {
public:
- explicit MockAddExpr(DataTypePtr type) : VExpr(type, false),
_type(std::move(type)) {}
+ explicit MockAddExpr(DataTypePtr type, std::vector<size_t>*
observed_batch_sizes = nullptr)
+ : VExpr(type, false),
+ _type(std::move(type)),
+ _observed_batch_sizes(observed_batch_sizes) {}
const std::string& expr_name() const override { return _name; }
Status execute_column_impl(VExprContext* context, const Block* block,
const Selector* selector,
size_t count, ColumnPtr& result_column) const
override {
+ if (_observed_batch_sizes != nullptr) {
+ _observed_batch_sizes->push_back(count);
+ }
ColumnPtr left;
ColumnPtr right;
RETURN_IF_ERROR(get_child(0)->execute_column(context, block, selector,
count, left));
@@ -187,9 +356,63 @@ private:
}
DataTypePtr _type;
+ std::vector<size_t>* _observed_batch_sizes;
std::string _name = "mock_add";
};
+class MockSparseCapturedInputExpr final : public VExpr {
+public:
+ MockSparseCapturedInputExpr(DataTypePtr type, size_t
captured_column_position,
+ bool* sparse_slots_use_column_nothing)
+ : VExpr(type, false),
+ _type(std::move(type)),
+ _captured_column_position(captured_column_position),
+
_sparse_slots_use_column_nothing(sparse_slots_use_column_nothing) {}
+
+ const std::string& expr_name() const override { return _name; }
+
+ Status execute_column_impl(VExprContext* context, const Block* block,
const Selector* selector,
+ size_t count, ColumnPtr& result_column) const
override {
+ *_sparse_slots_use_column_nothing = true;
+ for (size_t i = 0; i < _captured_column_position; ++i) {
+ if
(!check_and_get_column<ColumnNothing>(block->get_by_position(i).column.get())) {
+ *_sparse_slots_use_column_nothing = false;
+ break;
+ }
+ }
+
+ ColumnPtr captured;
+ ColumnPtr input;
+ RETURN_IF_ERROR(get_child(0)->execute_column(context, block, selector,
count, captured));
+ RETURN_IF_ERROR(get_child(1)->execute_column(context, block, selector,
count, input));
+ const IColumn* captured_data = captured.get();
+ if (const auto* nullable =
check_and_get_column<ColumnNullable>(captured_data)) {
+ captured_data = &nullable->get_nested_column();
+ }
+ const IColumn* input_data = input.get();
+ if (const auto* nullable =
check_and_get_column<ColumnNullable>(input_data)) {
+ input_data = &nullable->get_nested_column();
+ }
+ const auto& captured_values = assert_cast<const
ColumnInt32&>(*captured_data);
+ const auto& input_values = assert_cast<const
ColumnInt32&>(*input_data);
+ auto result = ColumnInt32::create();
+ result->reserve(count);
+ for (size_t i = 0; i < count; ++i) {
+ result->insert_value(captured_values.get_element(i) +
input_values.get_element(i));
+ }
+ result_column = std::move(result);
+ return Status::OK();
+ }
+
+ DataTypePtr execute_type(const Block* /*block*/) const override { return
_type; }
+
+private:
+ DataTypePtr _type;
+ size_t _captured_column_position;
+ bool* _sparse_slots_use_column_nothing;
+ std::string _name = "mock_sparse_captured_input";
+};
+
class MockMultiplyExpr final : public VExpr {
public:
explicit MockMultiplyExpr(DataTypePtr type) : VExpr(type, false),
_type(std::move(type)) {}
@@ -425,6 +648,460 @@ static void open_expr(const VExprSPtr& expr,
VExprContext* context) {
ASSERT_TRUE(expr->open(&state, context,
FunctionContext::THREAD_LOCAL).ok());
}
+static void open_expr_with_batch_size(const VExprSPtr& expr, VExprContext*
context,
+ int batch_size) {
+ RuntimeState state;
+ TQueryOptions query_options;
+ query_options.__set_batch_size(batch_size);
+ state.set_query_options(query_options);
+ RowDescriptor row_desc;
+ ASSERT_TRUE(expr->prepare(&state, row_desc, context).ok());
+ ASSERT_TRUE(expr->open(&state, context,
FunctionContext::THREAD_LOCAL).ok());
+}
+
+static void open_expr_with_block_budget(const VExprSPtr& expr, VExprContext*
context,
+ int batch_size, int64_t
preferred_block_size_bytes) {
+ RuntimeState state;
+ TQueryOptions query_options;
+ query_options.__set_batch_size(batch_size);
+ query_options.__set_preferred_block_size_bytes(preferred_block_size_bytes);
+ state.set_query_options(query_options);
+ RowDescriptor row_desc;
+ ASSERT_TRUE(expr->prepare(&state, row_desc, context).ok());
+ ASSERT_TRUE(expr->open(&state, context,
FunctionContext::THREAD_LOCAL).ok());
+}
+
+static const ColumnInt32& get_int_array_values(const ColumnPtr& result) {
+ const auto& result_array = assert_cast<const ColumnArray&>(*result);
+ const auto& nullable_values = assert_cast<const
ColumnNullable&>(*result_array.get_data_ptr());
+ return assert_cast<const
ColumnInt32&>(nullable_values.get_nested_column());
+}
+
+TEST(ArrayMapFunctionTest, IdentityLambdaSharesNestedInputColumn) {
+ auto int_type = std::make_shared<DataTypeInt32>();
+ auto array_int_type = std::make_shared<DataTypeArray>(int_type);
+ auto input = make_int_array_column({{1, 2}, {3}});
+ const auto& input_array = assert_cast<const ColumnArray&>(*input);
+ const IColumn* nested_input = input_array.get_data_ptr().get();
+
+ auto root =
VLambdaFunctionCallExpr::create_shared(make_lambda_call_node(array_int_type,
2));
+ auto lambda =
VLambdaFunctionExpr::create_shared(make_lambda_expr_node(int_type, {"x"}));
+ lambda->add_child(VColumnRef::create_shared(make_column_ref_node(0, "x",
int_type)));
+ root->add_child(lambda);
+ root->add_child(
+ std::make_shared<MockColumnExpr>(std::move(input), array_int_type,
"input_array"));
+
+ VExprContext context(root);
+ open_expr(root, &context);
+
+ Block block;
+ ColumnPtr result;
+ auto status = root->execute_column(&context, &block, nullptr, 2, result);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+
+ const auto& result_array = assert_cast<const ColumnArray&>(*result);
+ EXPECT_EQ(result_array.get_data_ptr().get(), nested_input);
+}
+
+TEST(ArrayMapFunctionTest, LambdaWithConstantCaptureUsesNestedColumnDirectly) {
+ auto int_type = std::make_shared<DataTypeInt32>();
+ auto array_int_type = std::make_shared<DataTypeArray>(int_type);
+ auto input = make_int_array_column({{1, 2}, {3}});
+ const auto& input_array = assert_cast<const ColumnArray&>(*input);
+ const IColumn* nested_input = input_array.get_data_ptr().get();
+
+ auto root =
VLambdaFunctionCallExpr::create_shared(make_lambda_call_node(array_int_type,
2));
+ auto lambda =
VLambdaFunctionExpr::create_shared(make_lambda_expr_node(int_type, {"x"}));
+ bool used_direct_input = false;
+ bool kept_capture_const = false;
+ auto body = std::make_shared<MockCapturedInputExpr>(int_type,
nested_input, &used_direct_input,
+ &kept_capture_const);
+ body->add_child(make_slot_ref(0, "captured", int_type));
+ body->add_child(VColumnRef::create_shared(make_column_ref_node(0, "x",
int_type)));
+ lambda->add_child(body);
+ root->add_child(lambda);
+ root->add_child(
+ std::make_shared<MockColumnExpr>(std::move(input), array_int_type,
"input_array"));
+
+ VExprContext context(root);
+ open_expr(root, &context);
+
+ Block block;
+ block.insert({ColumnConst::create(make_int_column({10}), 2), int_type,
"captured"});
+
+ ColumnPtr result;
+ auto status = root->execute_column(&context, &block, nullptr,
block.rows(), result);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ EXPECT_TRUE(used_direct_input);
+ EXPECT_TRUE(kept_capture_const);
+
+ const auto& values = get_int_array_values(result);
+ ASSERT_EQ(values.size(), 3);
+ EXPECT_EQ(values.get_element(0), 11);
+ EXPECT_EQ(values.get_element(1), 12);
+ EXPECT_EQ(values.get_element(2), 13);
+}
+
+TEST(ArrayMapFunctionTest, LargeLambdaProducesCorrectResult) {
+ auto int_type = std::make_shared<DataTypeInt32>();
+ auto array_int_type = std::make_shared<DataTypeArray>(int_type);
+ std::vector<int32_t> input_values(100000);
+ for (size_t i = 0; i < input_values.size(); ++i) {
+ input_values[i] = static_cast<int32_t>(i);
+ }
+
+ auto root =
VLambdaFunctionCallExpr::create_shared(make_lambda_call_node(array_int_type,
2));
+ auto lambda =
VLambdaFunctionExpr::create_shared(make_lambda_expr_node(int_type, {"x"}));
+ auto body = std::make_shared<MockIncrementExpr>(int_type);
+ body->add_child(VColumnRef::create_shared(make_column_ref_node(0, "x",
int_type)));
+ lambda->add_child(body);
+ root->add_child(lambda);
+ root->add_child(std::make_shared<MockColumnExpr>(
+ make_int_array_column({std::move(input_values)}), array_int_type,
"input_array"));
+
+ VExprContext context(root);
+ open_expr(root, &context);
+
+ Block block;
+ ColumnPtr result;
+ auto status = root->execute_column(&context, &block, nullptr, 1, result);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+
+ const auto& values = get_int_array_values(result);
+ ASSERT_EQ(values.size(), 100000);
+ EXPECT_EQ(values.get_element(0), 1);
+ EXPECT_EQ(values.get_element(99999), 100000);
+}
+
+TEST(ArrayMapFunctionTest, VariableLengthCaptureUsesOuterBatchSize) {
+ auto int_type = std::make_shared<DataTypeInt32>();
+ auto string_type = std::make_shared<DataTypeString>();
+ auto array_int_type = std::make_shared<DataTypeArray>(int_type);
+ std::vector<size_t> observed_batch_sizes;
+
+ auto root =
VLambdaFunctionCallExpr::create_shared(make_lambda_call_node(array_int_type,
2));
+ auto lambda =
VLambdaFunctionExpr::create_shared(make_lambda_expr_node(int_type, {"x"}));
+ auto body = std::make_shared<MockBatchSizeExpr>(int_type,
&observed_batch_sizes);
+ body->add_child(make_slot_ref(0, "captured", string_type));
+ lambda->add_child(body);
+ root->add_child(lambda);
+
root->add_child(std::make_shared<MockColumnExpr>(make_int_array_column({{1, 2,
3, 4, 5}}),
+ array_int_type,
"input_array"));
+
+ VExprContext context(root);
+ open_expr_with_batch_size(root, &context, 2);
+
+ auto captured = ColumnString::create();
+ std::string captured_value(5000, 'a');
+ captured->insert_data(captured_value.data(), captured_value.size());
+ Block block;
+ block.insert({std::move(captured), string_type, "captured"});
+
+ ColumnPtr result;
+ auto status = root->execute_column(&context, &block, nullptr,
block.rows(), result);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ ASSERT_EQ(observed_batch_sizes.size(), 3);
+ EXPECT_EQ(observed_batch_sizes[0], 2);
+ EXPECT_EQ(observed_batch_sizes[1], 2);
+ EXPECT_EQ(observed_batch_sizes[2], 1);
+ EXPECT_EQ(get_int_array_values(result).size(), 5);
+}
+
+TEST(ArrayMapFunctionTest, ComplexLambdaUsesOuterBatchSize) {
+ constexpr size_t intermediate_count = 100;
+ constexpr size_t nested_count = 1000;
+ constexpr int outer_batch_size = 256;
+ auto int_type = std::make_shared<DataTypeInt32>();
+ auto array_int_type = std::make_shared<DataTypeArray>(int_type);
+ std::vector<size_t> observed_batch_sizes;
+ std::vector<int32_t> input_values(nested_count);
+ for (size_t i = 0; i < nested_count; ++i) {
+ input_values[i] = static_cast<int32_t>(i);
+ }
+
+ auto root =
VLambdaFunctionCallExpr::create_shared(make_lambda_call_node(array_int_type,
2));
+ auto lambda =
VLambdaFunctionExpr::create_shared(make_lambda_expr_node(int_type, {"x"}));
+ auto body = std::make_shared<MockGreatestExpr>(int_type,
&observed_batch_sizes);
+ for (size_t i = 0; i < intermediate_count; ++i) {
+ auto increment = std::make_shared<MockIncrementExpr>(int_type);
+ increment->add_child(VColumnRef::create_shared(make_column_ref_node(0,
"x", int_type)));
+ body->add_child(increment);
+ }
+ lambda->add_child(body);
+ root->add_child(lambda);
+ root->add_child(std::make_shared<MockColumnExpr>(
+ make_int_array_column({std::move(input_values)}), array_int_type,
"input_array"));
+
+ VExprContext context(root);
+ open_expr_with_batch_size(root, &context, outer_batch_size);
+
+ Block block;
+ ColumnPtr result;
+ auto status = root->execute_column(&context, &block, nullptr, 1, result);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+
+ size_t observed_rows = 0;
+ for (size_t batch_rows : observed_batch_sizes) {
+ EXPECT_EQ(batch_rows,
+ std::min(static_cast<size_t>(outer_batch_size), nested_count
- observed_rows));
+ observed_rows += batch_rows;
+ }
+ EXPECT_EQ(observed_rows, nested_count);
+ EXPECT_EQ(observed_batch_sizes.size(), 4);
+
+ const auto& values = get_int_array_values(result);
+ ASSERT_EQ(values.size(), nested_count);
+ EXPECT_EQ(values.get_element(0), 1);
+ EXPECT_EQ(values.get_element(nested_count - 1), nested_count);
+}
+
+TEST(ArrayMapFunctionTest, FixedLengthInputsUseOuterBatchSize) {
+ constexpr size_t capture_count = 64;
+ constexpr size_t nested_count = 1000;
+ constexpr int outer_batch_size = 128;
+ auto int_type = std::make_shared<DataTypeInt32>();
+ auto int64_type = std::make_shared<DataTypeInt64>();
+ auto array_int_type = std::make_shared<DataTypeArray>(int_type);
+ std::vector<size_t> observed_batch_sizes;
+
+ auto root =
VLambdaFunctionCallExpr::create_shared(make_lambda_call_node(array_int_type,
2));
+ auto lambda =
VLambdaFunctionExpr::create_shared(make_lambda_expr_node(int_type, {"x"}));
+ auto body = std::make_shared<MockBatchSizeExpr>(int_type,
&observed_batch_sizes);
+ for (size_t i = 0; i < capture_count; ++i) {
+ body->add_child(
+ make_slot_ref(static_cast<int>(i), "captured_" +
std::to_string(i), int64_type));
+ }
+ lambda->add_child(body);
+ root->add_child(lambda);
+ root->add_child(std::make_shared<MockColumnExpr>(
+ make_int_array_column({std::vector<int32_t>(nested_count, 1)}),
array_int_type,
+ "input_array"));
+
+ VExprContext context(root);
+ open_expr_with_batch_size(root, &context, outer_batch_size);
+
+ Block block;
+ for (size_t i = 0; i < capture_count; ++i) {
+ auto captured = ColumnInt64::create();
+ captured->insert_value(static_cast<int64_t>(i));
+ block.insert({std::move(captured), int64_type, "captured_" +
std::to_string(i)});
+ }
+
+ ColumnPtr result;
+ auto status = root->execute_column(&context, &block, nullptr,
block.rows(), result);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+
+ size_t observed_rows = 0;
+ for (size_t batch_rows : observed_batch_sizes) {
+ EXPECT_EQ(batch_rows,
+ std::min(static_cast<size_t>(outer_batch_size), nested_count
- observed_rows));
+ observed_rows += batch_rows;
+ }
+ EXPECT_EQ(observed_rows, nested_count);
+ EXPECT_EQ(observed_batch_sizes.size(), 8);
+ EXPECT_EQ(get_int_array_values(result).size(), nested_count);
+}
+
+TEST(ArrayMapFunctionTest, FixedLengthInputsUsePreferredBlockSizeBudget) {
+ constexpr size_t capture_count = 64;
+ constexpr size_t nested_count = 3000;
+ constexpr int outer_batch_size = 65535;
+ constexpr int64_t preferred_block_size_bytes = 1024 * 1024;
+ auto int_type = std::make_shared<DataTypeInt32>();
+ auto int64_type = std::make_shared<DataTypeInt64>();
+ auto array_int_type = std::make_shared<DataTypeArray>(int_type);
+ std::vector<size_t> observed_batch_sizes;
+
+ const bool old_enable_adaptive_batch_size =
config::enable_adaptive_batch_size;
+ config::enable_adaptive_batch_size = true;
+ Defer restore_adaptive_batch_size {[old_enable_adaptive_batch_size] {
+ config::enable_adaptive_batch_size = old_enable_adaptive_batch_size;
+ }};
+
+ auto root =
VLambdaFunctionCallExpr::create_shared(make_lambda_call_node(array_int_type,
2));
+ auto lambda =
VLambdaFunctionExpr::create_shared(make_lambda_expr_node(int_type, {"x"}));
+ auto body = std::make_shared<MockBatchSizeExpr>(int_type,
&observed_batch_sizes);
+ for (size_t i = 0; i < capture_count; ++i) {
+ body->add_child(
+ make_slot_ref(static_cast<int>(i), "captured_" +
std::to_string(i), int64_type));
+ }
+ lambda->add_child(body);
+ root->add_child(lambda);
+ root->add_child(std::make_shared<MockColumnExpr>(
+ make_int_array_column({std::vector<int32_t>(nested_count, 1)}),
array_int_type,
+ "input_array"));
+
+ VExprContext context(root);
+ open_expr_with_block_budget(root, &context, outer_batch_size,
preferred_block_size_bytes);
+
+ Block block;
+ for (size_t i = 0; i < capture_count; ++i) {
+ auto captured = ColumnInt64::create();
+ captured->insert_value(static_cast<int64_t>(i));
+ block.insert({std::move(captured), int64_type, "captured_" +
std::to_string(i)});
+ }
+
+ ColumnPtr result;
+ auto status = root->execute_column(&context, &block, nullptr,
block.rows(), result);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+
+ ASSERT_GT(observed_batch_sizes.size(), 1);
+ size_t observed_rows = 0;
+ for (size_t batch_rows : observed_batch_sizes) {
+ EXPECT_LE(batch_rows, preferred_block_size_bytes / (capture_count *
sizeof(int64_t)));
+ observed_rows += batch_rows;
+ }
+ EXPECT_EQ(observed_rows, nested_count);
+ EXPECT_EQ(get_int_array_values(result).size(), nested_count);
+}
+
+TEST(ArrayMapFunctionTest,
MultiBatchPreservesCaptureMappingAcrossSelectedArrayRows) {
+ constexpr size_t selected_array_size = 300;
+ constexpr int outer_batch_size = 400;
+ auto int_type = std::make_shared<DataTypeInt32>();
+ auto array_int_type = std::make_shared<DataTypeArray>(int_type);
+ std::vector<size_t> observed_batch_sizes;
+
+ auto root =
VLambdaFunctionCallExpr::create_shared(make_lambda_call_node(array_int_type,
2));
+ auto lambda =
VLambdaFunctionExpr::create_shared(make_lambda_expr_node(int_type, {"x"}));
+ auto body = std::make_shared<MockAddExpr>(int_type, &observed_batch_sizes);
+ body->add_child(make_slot_ref(0, "captured_0", int_type));
+ body->add_child(VColumnRef::create_shared(make_column_ref_node(0, "x",
int_type)));
+ lambda->add_child(body);
+ root->add_child(lambda);
+ root->add_child(make_slot_ref(1, "input_array", array_int_type));
+
+ VExprContext context(root);
+ open_expr_with_batch_size(root, &context, outer_batch_size);
+
+ std::vector<std::vector<int32_t>> input_rows(6);
+ input_rows[0] = {-1};
+ input_rows[1].resize(selected_array_size);
+ for (size_t i = 0; i < selected_array_size; ++i) {
+ input_rows[1][i] = static_cast<int32_t>(i);
+ }
+ input_rows[2] = {-2, -3};
+ input_rows[4] = {-4};
+ input_rows[5].resize(selected_array_size);
+ for (size_t i = 0; i < selected_array_size; ++i) {
+ input_rows[5][i] = 1000 + static_cast<int32_t>(i);
+ }
+
+ Block block;
+ block.insert({make_int_column({100, 10, 200, 30, 300, 50}), int_type,
"captured_0"});
+ block.insert({make_int_array_column(input_rows), array_int_type,
"input_array"});
+
+ Selector selector;
+ selector.push_back(1);
+ selector.push_back(3);
+ selector.push_back(5);
+
+ ColumnPtr result;
+ auto status = root->execute_column(&context, &block, &selector,
selector.size(), result);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+
+ const size_t total_nested_rows = 2 * selected_array_size;
+ ASSERT_EQ(observed_batch_sizes.size(), 2);
+ EXPECT_EQ(observed_batch_sizes[0], outer_batch_size);
+ EXPECT_EQ(observed_batch_sizes[1], total_nested_rows - outer_batch_size);
+
+ const auto& result_array = assert_cast<const ColumnArray&>(*result);
+ ASSERT_EQ(result_array.size(), 3);
+ EXPECT_EQ(result_array.get_offsets()[0], selected_array_size);
+ EXPECT_EQ(result_array.get_offsets()[1], selected_array_size);
+ EXPECT_EQ(result_array.get_offsets()[2], total_nested_rows);
+
+ const auto& values = get_int_array_values(result);
+ ASSERT_EQ(values.size(), total_nested_rows);
+ EXPECT_EQ(values.get_element(0), 10);
+ EXPECT_EQ(values.get_element(selected_array_size - 1), 309);
+ EXPECT_EQ(values.get_element(selected_array_size), 1050);
+ EXPECT_EQ(values.get_element(outer_batch_size - 1),
+ 50 + 1000 + static_cast<int32_t>(outer_batch_size -
selected_array_size - 1));
+ EXPECT_EQ(values.get_element(outer_batch_size),
+ 50 + 1000 + static_cast<int32_t>(outer_batch_size -
selected_array_size));
+ EXPECT_EQ(values.get_element(total_nested_rows - 1), 1349);
+}
+
+TEST(ArrayMapFunctionTest,
SparseCapturedColumnUsesColumnNothingForUnusedSlots) {
+ auto int_type = std::make_shared<DataTypeInt32>();
+ auto uint8_type = std::make_shared<DataTypeUInt8>();
+ auto array_int_type = std::make_shared<DataTypeArray>(int_type);
+ constexpr size_t captured_column_position = 999;
+ std::vector<int32_t> input_values(1000, 1);
+
+ auto root =
VLambdaFunctionCallExpr::create_shared(make_lambda_call_node(array_int_type,
2));
+ auto lambda =
VLambdaFunctionExpr::create_shared(make_lambda_expr_node(int_type, {"x"}));
+ bool sparse_slots_use_column_nothing = false;
+ auto body = std::make_shared<MockSparseCapturedInputExpr>(int_type,
captured_column_position,
+
&sparse_slots_use_column_nothing);
+ body->add_child(make_slot_ref(captured_column_position, "captured",
int_type));
+ body->add_child(VColumnRef::create_shared(make_column_ref_node(0, "x",
int_type)));
+ lambda->add_child(body);
+ root->add_child(lambda);
+ root->add_child(std::make_shared<MockColumnExpr>(
+ make_int_array_column({std::move(input_values)}), array_int_type,
"input_array"));
+
+ VExprContext context(root);
+ open_expr(root, &context);
+
+ Block block;
+ for (size_t i = 0; i < captured_column_position; ++i) {
+ block.insert({ColumnUInt8::create(1, 0), uint8_type, "unused"});
+ }
+ block.insert({make_int_column({10}), int_type, "captured"});
+
+ ColumnPtr result;
+ auto status = root->execute_column(&context, &block, nullptr, 1, result);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ EXPECT_TRUE(sparse_slots_use_column_nothing);
+
+ const auto& values = get_int_array_values(result);
+ ASSERT_EQ(values.size(), 1000);
+ EXPECT_EQ(values.get_element(0), 11);
+ EXPECT_EQ(values.get_element(999), 11);
+}
+
+TEST(ArrayMapFunctionTest, CapturedColumnExpansionUsesOuterSelector) {
+ auto int_type = std::make_shared<DataTypeInt32>();
+ auto array_int_type = std::make_shared<DataTypeArray>(int_type);
+
+ auto root =
VLambdaFunctionCallExpr::create_shared(make_lambda_call_node(array_int_type,
2));
+ auto lambda =
VLambdaFunctionExpr::create_shared(make_lambda_expr_node(int_type, {"x"}));
+ auto add = std::make_shared<MockAddExpr>(int_type);
+ add->add_child(make_slot_ref(0, "captured", int_type));
+ add->add_child(VColumnRef::create_shared(make_column_ref_node(0, "x",
int_type)));
+ lambda->add_child(add);
+ root->add_child(lambda);
+ root->add_child(make_slot_ref(1, "input_array", array_int_type));
+
+ VExprContext context(root);
+ open_expr(root, &context);
+
+ Block block;
+ block.insert({make_int_column({10, 20, 30, 40}), int_type, "captured"});
+ block.insert(
+ {make_int_array_column({{1, 2}, {3}, {}, {7, 8, 9}}),
array_int_type, "input_array"});
+ Selector selector;
+ selector.push_back(1);
+ selector.push_back(3);
+
+ ColumnPtr result;
+ auto status = root->execute_column(&context, &block, &selector,
selector.size(), result);
+ ASSERT_TRUE(status.ok()) << status.to_string();
+
+ const auto& result_array = assert_cast<const ColumnArray&>(*result);
+ ASSERT_EQ(result_array.size(), 2);
+ ASSERT_EQ(result_array.get_offsets()[0], 1);
+ ASSERT_EQ(result_array.get_offsets()[1], 4);
+ const auto& values = get_int_array_values(result);
+ ASSERT_EQ(values.size(), 4);
+ EXPECT_EQ(values.get_element(0), 23);
+ EXPECT_EQ(values.get_element(1), 47);
+ EXPECT_EQ(values.get_element(2), 48);
+ EXPECT_EQ(values.get_element(3), 49);
+}
+
TEST(ArrayMapFunctionTest, NestedLambdaWithSameArgumentNameUsesInnerScope) {
auto int_type = std::make_shared<DataTypeInt32>();
auto array_int_type = std::make_shared<DataTypeArray>(int_type);
diff --git
a/regression-test/data/query_p0/sql_functions/array_functions/test_array_map_function_with_column.out
b/regression-test/data/query_p0/sql_functions/array_functions/test_array_map_function_with_column.out
index 24d7b53e2b3..dce1c723407 100644
---
a/regression-test/data/query_p0/sql_functions/array_functions/test_array_map_function_with_column.out
+++
b/regression-test/data/query_p0/sql_functions/array_functions/test_array_map_function_with_column.out
@@ -39,3 +39,17 @@
4 5 [6, 7, null, 9] [4, 5, 6, 7] [0, 0, null, 0]
5 6 [10, 11, 12, 13] [8, 9, null, 11] [0, 0, null, 0]
+-- !select_captured_selector_if --
+0 []
+1 [23]
+2 []
+3 [47, 48, 49]
+4 [null, null]
+
+-- !select_captured_selector_case --
+0 [-9, -8]
+1 [23]
+2 []
+3 [47, 48, 49]
+4 [null, null]
+
diff --git
a/regression-test/suites/query_p0/sql_functions/array_functions/test_array_map_function_with_column.groovy
b/regression-test/suites/query_p0/sql_functions/array_functions/test_array_map_function_with_column.groovy
index 54904b64899..f5513be21b8 100644
---
a/regression-test/suites/query_p0/sql_functions/array_functions/test_array_map_function_with_column.groovy
+++
b/regression-test/suites/query_p0/sql_functions/array_functions/test_array_map_function_with_column.groovy
@@ -66,5 +66,30 @@ suite("test_array_map_function_with_column") {
qt_select_7 "select *,array_map((x,y)->x+k1+k2 >
y+k1*k2,c_array1,c_array2) from ${tableName} where array_count((x,y) ->
k1*x>y+k2, c_array1, c_array2) > 1 order by k1;"
+ sql "truncate table ${tableName};"
+ sql """INSERT INTO ${tableName} values
+ (0, 10, [1,2], [1,2]),
+ (1, 20, [3], [3]),
+ (2, 30, [], []),
+ (3, 40, [7,8,9], [7,8,9]),
+ (4, NULL, [10,NULL], [10,NULL]);
+ """
+
+ sql "set short_circuit_evaluation = true;"
+ qt_select_captured_selector_if """
+ select k1, if(k1 in (1, 3, 4), array_map(x -> x + k2, c_array1), [])
+ from ${tableName}
+ order by k1
+ """
+ qt_select_captured_selector_case """
+ select k1,
+ case when k1 in (1, 3, 4)
+ then array_map(x -> x + k2, c_array1)
+ else array_map(x -> x - k2, c_array1)
+ end
+ from ${tableName}
+ order by k1
+ """
+
// sql "DROP TABLE IF EXISTS ${tableName}"
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
