yujun777 commented on code in PR #69456: URL: https://github.com/apache/doris/pull/69456#discussion_r4356686316
########## be/src/exprs/lambda_function/varray_sorted_merge_function.cpp: ########## @@ -0,0 +1,215 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include <memory> +#include <string> +#include <utility> + +#include "common/status.h" +#include "core/assert_cast.h" +#include "core/block/block.h" +#include "core/block/column_with_type_and_name.h" +#include "core/column/column.h" +#include "core/column/column_array.h" +#include "core/column/column_nullable.h" +#include "core/column/column_vector.h" +#include "core/data_type/data_type.h" +#include "core/data_type/data_type_array.h" +#include "core/data_type/data_type_nullable.h" +#include "exprs/lambda_function/lambda_execution_context.h" +#include "exprs/lambda_function/lambda_function.h" +#include "exprs/lambda_function/lambda_function_factory.h" +#include "exprs/vexpr.h" +#include "exprs/vexpr_context.h" +#include "exprs/vlambda_function_expr.h" + +namespace doris { + +/** + * array_sorted_merge(merge_fn, arr1, arr2, is_asc) + * + * Merges two arrays that are already sorted by the elements' natural order. The merge lambda + * receives the two current candidates and returns the combined element; returning NULL means the + * two elements must not be combined, so both are kept. Each input key occurs at most once per + * array, so a key appears at most twice across the two arrays and the lambda is called at most + * once per key. + */ +class ArraySortedMergeFunction : public LambdaFunction { + ENABLE_FACTORY_CREATOR(ArraySortedMergeFunction); + +public: + ~ArraySortedMergeFunction() override = default; + + static constexpr auto name = "array_sorted_merge"; + + static LambdaFunctionPtr create() { return std::make_shared<ArraySortedMergeFunction>(); } + + std::string get_name() const override { return name; } + + Status execute(VExprContext* context, const Block* block, const Selector* expr_selector, + size_t count, ColumnPtr& result_column, const DataTypePtr& result_type, + const VExprSPtrs& children) const override { + DCHECK_EQ(children.size(), 4); + + const ColumnArray* arrays[2] = {nullptr, nullptr}; + const IColumn* data[2] = {nullptr, nullptr}; + DataTypePtr nested_types[2]; + ColumnPtr holders[2]; + const ColumnUInt8::Container* nullmaps[2] = {nullptr, nullptr}; + for (int i = 0; i < 2; ++i) { + RETURN_IF_ERROR(children[i + 1]->execute_column(context, block, expr_selector, count, + holders[i])); + holders[i] = holders[i]->convert_to_full_column_if_const(); + DataTypePtr type = children[i + 1]->execute_type(block); + const IColumn* array_column = holders[i].get(); + if (const auto* nullable = check_and_get_column<ColumnNullable>(array_column)) { + nullmaps[i] = &nullable->get_null_map_column().get_data(); + array_column = &nullable->get_nested_column(); + type = assert_cast<const DataTypeNullable&>(*type).get_nested_type(); + } + arrays[i] = assert_cast<const ColumnArray*>(array_column); + data[i] = arrays[i]->get_data_ptr().get(); + nested_types[i] = assert_cast<const DataTypeArray&>(*type).get_nested_type(); + } + + ColumnPtr flag_holder; + RETURN_IF_ERROR( + children[3]->execute_column(context, block, expr_selector, count, flag_holder)); + flag_holder = flag_holder->convert_to_full_column_if_const(); + const IColumn* flag_column = flag_holder.get(); + const ColumnUInt8::Container* flag_nullmap = nullptr; + if (const auto* nullable = check_and_get_column<ColumnNullable>(flag_column)) { + flag_nullmap = &nullable->get_null_map_column().get_data(); + flag_column = &nullable->get_nested_column(); + } + const auto& flags = assert_cast<const ColumnUInt8&>(*flag_column).get_data(); + + const size_t rows = arrays[0]->size(); + auto result_nested = nested_types[0]->create_column(); + auto result_offsets = ColumnArray::ColumnOffsets::create(); + auto& offsets_data = result_offsets->get_data(); + offsets_data.reserve(rows); + auto result_nullmap = ColumnUInt8::create(rows, 0); + auto& result_nullmap_data = result_nullmap->get_data(); + + // Two one-row columns hold the lambda arguments of the current merge step. They carry + // the lambda argument names because the column refs resolve by name when no + // position-based gap is set. + const auto* lambda_expr = assert_cast<const VLambdaFunctionExpr*>(children[0].get()); + const std::vector<std::string>* argument_names = + lambda_expr->has_argument_names() ? &lambda_expr->argument_names() : nullptr; + Block lambda_block; + for (int i = 0; i < 2; ++i) { + const std::string argument_name = + argument_names != nullptr && i < argument_names->size() + ? (*argument_names)[i] + : "merge_arg_" + std::to_string(i); + lambda_block.insert(ColumnWithTypeAndName(nested_types[i]->create_column(), + nested_types[i], argument_name)); + } + // The lambda arguments are bound by position (column ids 0 and 1), like array_sort's + // comparator, so the body reads these two one-row columns. + IColumn* argument_columns[2] = { + lambda_block.get_by_position(0).column->assert_mutable().get(), + lambda_block.get_by_position(1).column->assert_mutable().get()}; + + LambdaExecutionContext::Frame lambda_frame; + lambda_frame.bind_by_name = false; + lambda_frame.parent_bindings_visible = false; + LambdaExecutionContext::FrameGuard lambda_frame_guard(context->lambda_execution_context(), + std::move(lambda_frame)); + + const auto& offsets0 = arrays[0]->get_offsets(); + const auto& offsets1 = arrays[1]->get_offsets(); + const int lambda_result_base = static_cast<int>(lambda_block.columns()); + + for (size_t row = 0; row < rows; ++row) { + const bool row_is_null = (nullmaps[0] != nullptr && (*nullmaps[0])[row] != 0) || + (nullmaps[1] != nullptr && (*nullmaps[1])[row] != 0); + if (row_is_null) { + result_nullmap_data[row] = 1; + offsets_data.push_back(result_nested->size()); + continue; + } + if (flag_nullmap != nullptr && (*flag_nullmap)[row] != 0) { + return Status::InvalidArgument( + "array_sorted_merge requires a non-null ascending flag"); + } + const bool ascending = flags[row] != 0; + + size_t left = row == 0 ? 0 : offsets0[row - 1]; + const size_t left_end = offsets0[row]; + size_t right = row == 0 ? 0 : offsets1[row - 1]; + const size_t right_end = offsets1[row]; + + while (left < left_end || right < right_end) { + if (left == left_end) { + result_nested->insert_from(*data[1], right++); + continue; + } + if (right == right_end) { + result_nested->insert_from(*data[0], left++); + continue; + } + + argument_columns[0]->clear(); + argument_columns[0]->insert_from(*data[0], left); + argument_columns[1]->clear(); + argument_columns[1]->insert_from(*data[1], right); + + int lambda_res_id = lambda_result_base; + RETURN_IF_ERROR(children[0]->execute(context, &lambda_block, &lambda_res_id)); Review Comment: Fixed in 6ac1f1aea13: an enclosing lambda's argument now resolves correctly. The function collects the named bindings visible at the current execution point (`LambdaExecutionContext::collect_visible_bindings`) and publishes their values under those names in its one-row merge block, before the two merge arguments so that a merge argument with the same name still shadows an outer one. The values are read at the same row as the arrays, so short-circuit evaluation and batched lambda execution keep reading the right value. `order_qt_capture_enclosing_lambda` in `test_array_sorted_merge` covers `array_map(x -> array_sorted_merge((a, b) -> IF(a = b AND a < x, a, NULL), [1], [1], true), [0, 2])` and returns `[[1, 1], [1]]` (x=0 keeps both values, x=2 merges them). A nested lambda that captures one of the merge arguments is still rejected during prepare, because those arguments are not published into nested frames. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
