github-actions[bot] commented on code in PR #69456: URL: https://github.com/apache/doris/pull/69456#discussion_r4384698012
########## 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: [P2] Materialize virtual captures before copying the lambda inputs With `experimental_enable_virtual_slot_for_cse=true`, a DUP table `t(lim INT NOT NULL, l ARRAY<INT>, r ARRAY<INT>)` containing `(2,[1],[1])` should return 2 for: ```sql SELECT lim FROM t WHERE abs(lim)>0 AND abs(lim)<10 AND cardinality(array_sorted_merge( (a,b)->IF(a=b AND a<abs(lim),a,NULL),l,r,true))=1; ``` `PushDownVirtualColumnsIntoOlapScan` extracts the repeated expression from the ordinary conjuncts, then replaces it inside the lambda too. The reduced rewritten tree inferred from that rule is: ```text Filter(v>0, v<10, cardinality(merge((a,b)->IF(a=b AND a<v,a,NULL),l,r))=1) OlapScan(t; virtual v=abs(lim)) ``` For a pushed scan filter, `SegmentIterator::_process_common_expr` runs before final virtual-column materialization, so v can still be `ColumnNothing`. `VirtualSlotRef::collect_slot_column_ids` records only v, excluding its hidden lim dependency. This copy therefore preserves the lazy placeholder, while the private block fills lim's position with `UInt8 temp`. When `VirtualSlotRef` lazily evaluates `abs(lim)`, it reads that wrong type and `FunctionUnaryArithmetic` returns an expected-data-type error instead of the matching row. Earlier virtual comparisons do not cache their result in the block. Materialize virtual captures against the original block before cutting them, or preserve their transitive dependencies and types, and add a pushed-filter CSE regression. This is distinct from the repaired ordinary-slot capture thread: the missing inputs belong to an out-of-tree virtual expression. -- 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]
