yujun777 commented on code in PR #69456: URL: https://github.com/apache/doris/pull/69456#discussion_r4384810800
########## be/src/exprs/lambda_function/varray_sorted_merge_function.cpp: ########## @@ -0,0 +1,377 @@ +// 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: a merge argument + /// stands for the candidate of one comparison step, which is not part of the supported + /// contract for a nested lambda. Nested lambdas may still use their own arguments + /// (shadowing) and may capture enclosing lambda arguments. + 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 has no defined meaning + // (the argument is the candidate of a single comparison step); reject it during prepare + // instead of resolving it against some column at runtime. + 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); + // An enclosing lambda's argument is a name reference rather than a slot, and this + // function executes its lambda against a block of its own that hides the parent frames. + // Publish the visible binding values under their names in that block, so the reference + // resolves by name like the merge arguments do. + std::vector<LambdaExecutionContext::Binding> enclosing_bindings; + context->lambda_execution_context().collect_visible_bindings(enclosing_bindings); + if ((!captured_ids.empty() || !enclosing_bindings.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(); Review Comment: Fixed in 5309c29db2b: the private merge block now also copies the columns a captured virtual slot's expression reads. `VirtualSlotRef::collect_slot_column_ids` only reports the virtual column itself, so before this the block kept the `ColumnNothing` placeholder for `v` and padded `lim`'s position with the `UInt8 temp` column; the lazy `abs(lim)` then read a wrong-typed column. The new `collect_virtual_capture_inputs` walks the lambda for `VirtualSlotRef` nodes, adds the slot ids of their virtual expressions recursively, and those columns are copied at their original ids with their real types. Reproduced before the fix with the reduced query and `experimental_enable_virtual_slot_for_cse = true`: `abs's argument does not match the expected data type.` The same query now returns `2`, and `order_qt_virtual_column_capture` in `test_array_sorted_merge` pins the pushed-filter CSE case (DUP table `t(lim, l, r)` with `(2,[1],[1])`). Note this covers `array_sorted_merge`; a general fix could instead make `VirtualSlotRef::collect_slot_column_ids` include its expression's inputs so every higher-order function gets the same behavior, but that changes the contract for all callers, so it is left for a separate change. -- 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]
