yujun777 commented on code in PR #69456: URL: https://github.com/apache/doris/pull/69456#discussion_r4355468827
########## be/src/exprs/lambda_function/varray_sorted_merge_function.cpp: ########## @@ -0,0 +1,323 @@ +// 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 <algorithm> +#include <map> +#include <memory> +#include <set> +#include <string> +#include <utility> +#include <vector> + +#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 "core/data_type/data_type_number.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; } + + /// Rejects a nested lambda that captures one of this lambda's arguments: the position-based + /// frame below does not publish them into nested lambdas, so such a capture would resolve + /// against the wrong column. Nested lambdas may still use their own arguments (shadowing). + static Status validate_nested_captures(const VExprSPtr& expr, + const std::vector<std::string>& argument_names, + bool inside_nested_lambda, + std::vector<std::string> shadowed_names) { + if (expr->is_column_ref()) { + const std::string& column_name = expr->expr_name(); + if (std::ranges::find(shadowed_names, column_name) != shadowed_names.end()) { + return Status::OK(); + } + if (inside_nested_lambda && + std::ranges::find(argument_names, column_name) != argument_names.end()) { + return Status::NotSupported( + "array_sorted_merge merge lambda argument '{}' cannot be captured by a " + "nested lambda", + column_name); + } + return Status::OK(); + } + if (expr->node_type() == TExprNodeType::LAMBDA_FUNCTION_EXPR) { + const auto* lambda = assert_cast<const VLambdaFunctionExpr*>(expr.get()); + if (lambda->has_argument_names()) { + for (const auto& argument_name : lambda->argument_names()) { + shadowed_names.emplace_back(argument_name); + } + } + return validate_nested_captures(lambda->get_child(0), argument_names, true, + std::move(shadowed_names)); + } + for (const auto& child : expr->children()) { + RETURN_IF_ERROR(validate_nested_captures(child, argument_names, inside_nested_lambda, + shadowed_names)); + } + return Status::OK(); + } + + Status prepare(RuntimeState* state, const VExprSPtrs& children) override { + RETURN_IF_ERROR(LambdaFunction::prepare(state, children)); + DCHECK_EQ(children.size(), 4); + const auto* lambda_expr = assert_cast<const VLambdaFunctionExpr*>(children[0].get()); + if (!lambda_expr->has_argument_names()) { + return Status::OK(); + } + // A nested lambda that captures one of this lambda's arguments would need a binding + // that does not exist in the position-based frame below; reject it during prepare so + // the query fails instead of resolving the argument against the wrong column. + return validate_nested_captures(children[0]->get_child(0), lambda_expr->argument_names(), + false, {}); + } + + 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(); + + // The lambda may capture enclosing columns (for example a table column compared against + // the candidates). Materialize every captured input at its original position in the + // one-row lambda block, so slot refs read the current row; the two merge candidates + // take the positions after the last captured column and are bound by name. + std::set<int> captured_ids; + children[0]->collect_slot_column_ids(captured_ids); + if (!captured_ids.empty() && block == nullptr) { + return Status::InternalError( + "array_sorted_merge lambda references outer columns but has no input block"); + } + std::map<int, ColumnPtr> captured_columns; + for (int id : captured_ids) { + if (id < 0 || static_cast<size_t>(id) >= block->columns()) { + return Status::InternalError( + "array_sorted_merge lambda input column id {} is outside input block", id); + } + captured_columns[id] = + block->get_by_position(id).column->convert_to_full_column_if_const(); + } + const int lambda_argument_base = captured_ids.empty() ? 0 : *captured_ids.rbegin() + 1; + 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 < lambda_argument_base; ++i) { + auto captured = captured_columns.find(i); + if (captured != captured_columns.end()) { + lambda_block.insert(ColumnWithTypeAndName(captured->second->clone_empty(), + block->get_by_position(i).type, + block->get_by_position(i).name)); + } else { + lambda_block.insert(ColumnWithTypeAndName( + ColumnUInt8::create(1, 0), std::make_shared<DataTypeUInt8>(), "temp")); + } + } + 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 name (the column refs carry the argument names), + // so the body reads these two one-row columns. + IColumn* argument_columns[2] = { + lambda_block.get_by_position(lambda_argument_base).column->assert_mutable().get(), + lambda_block.get_by_position(lambda_argument_base + 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()); + + // The merged element may come from a non-nullable lambda body (for example COALESCE), + // while array elements are always nullable; normalize before inserting. + auto append_element = [&](const IColumn& source, size_t source_row) { + if (auto* nullable_target = check_and_get_column<ColumnNullable>(result_nested.get())) { + if (const auto* nullable_source = check_and_get_column<ColumnNullable>(&source)) { + nullable_target->insert_from(*nullable_source, source_row); + } else { + nullable_target->get_nested_column().insert_from(source, source_row); + nullable_target->get_null_map_data().push_back(0); + } + } else { + result_nested->insert_from(source, source_row); + } + }; + + 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; + for (const auto& [id, column] : captured_columns) { + lambda_block.get_by_position(id).column = column->cut(row, 1); Review Comment: Fixed in 49eb8ebc9c7: captured columns are now cut at `expr_selector[row]`, the same original row the array children read, so a short-circuit branch that only sees the selected rows reads the right captured value. `test_array_sorted_merge` adds `order_qt_short_circuit_selector`: with `short_circuit_evaluation = true`, row id=1 (lim=2) returns `[1]`; reading the compacted row 0 (lim=0) used to return `[1,1]`. -- 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]
