diff --git a/be/src/exprs/aggregate/aggregate_function_array_agg.cpp b/be/src/exprs/aggregate/aggregate_function_array_agg.cpp index 0fddadbaab291e..e5e7fb5670920e 100644 --- a/be/src/exprs/aggregate/aggregate_function_array_agg.cpp +++ b/be/src/exprs/aggregate/aggregate_function_array_agg.cpp @@ -20,6 +20,7 @@ #include "core/call_on_type_index.h" #include "exprs/aggregate/aggregate_function_collect.h" #include "exprs/aggregate/aggregate_function_simple_factory.h" +#include "exprs/aggregate/factory_helpers.h" #include "exprs/aggregate/helpers.h" namespace doris { @@ -60,7 +61,57 @@ AggregateFunctionPtr create_aggregate_function_array_agg(const std::string& name return agg_fn; } +// The element column of array_agg_if is either the declared nullable column or the raw column; +// NullableElem picks the matching Data specialization at create time so the per-row add() path +// never pays a runtime nullability check. +template +AggregateFunctionPtr do_create_array_agg_if(const DataTypes& argument_types, + const bool result_is_nullable, + const AggregateFunctionAttr& attr) { + // array_agg_if(cond, elem): the element type lives at argument index 1. + return creator_without_type::create_ignore_nullable< + AggregateFunctionArrayAggIf< + AggregateFunctionArrayAggData>>( + argument_types, result_is_nullable, attr); +} + +AggregateFunctionPtr create_aggregate_function_array_agg_if(const std::string& name, + const DataTypes& argument_types, + const DataTypePtr& result_type, + const bool result_is_nullable, + const AggregateFunctionAttr& attr) { + assert_arity_range(name, argument_types, 2, 2); + AggregateFunctionPtr agg_fn; + auto call = [&](const auto& type) -> bool { + using DispatcType = std::decay_t; + if (argument_types[1]->is_nullable()) { + agg_fn = do_create_array_agg_if( + argument_types, result_is_nullable, attr); + } else { + agg_fn = do_create_array_agg_if( + argument_types, result_is_nullable, attr); + } + return true; + }; + + if (!dispatch_switch_all(argument_types[1]->get_primitive_type(), call)) { + // We do not care what the real type is; complex element types still respect nullability. + if (argument_types[1]->is_nullable()) { + agg_fn = do_create_array_agg_if(argument_types, result_is_nullable, + attr); + } else { + agg_fn = do_create_array_agg_if(argument_types, result_is_nullable, + attr); + } + } + return agg_fn; +} + void register_aggregate_function_array_agg(AggregateFunctionSimpleFactory& factory) { factory.register_function_both("array_agg", create_aggregate_function_array_agg); } -} // namespace doris \ No newline at end of file + +void register_aggregate_function_array_agg_if(AggregateFunctionSimpleFactory& factory) { + factory.register_function_both("array_agg_if", create_aggregate_function_array_agg_if); +} +} // namespace doris diff --git a/be/src/exprs/aggregate/aggregate_function_array_agg.h b/be/src/exprs/aggregate/aggregate_function_array_agg.h index 002e11cdd13d5f..49a29e6806e065 100644 --- a/be/src/exprs/aggregate/aggregate_function_array_agg.h +++ b/be/src/exprs/aggregate/aggregate_function_array_agg.h @@ -24,6 +24,7 @@ #include "core/column/column_decimal.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.h" #include "core/data_type/data_type_array.h" #include "core/data_type/data_type_nullable.h" @@ -35,19 +36,19 @@ namespace doris { class Arena; -template +template struct AggregateFunctionArrayAggData { static constexpr PrimitiveType PType = T; static constexpr bool use_native_serde = false; using ElementType = typename PrimitiveTypeTraits::CppType; using ColVecType = typename PrimitiveTypeTraits::ColumnType; - using Self = AggregateFunctionArrayAggData; + using Self = AggregateFunctionArrayAggData; MutableColumnPtr column_data; ColVecType* nested_column = nullptr; NullMap* null_map = nullptr; AggregateFunctionArrayAggData(const DataTypes& argument_types) { - DataTypePtr column_type = make_nullable(argument_types[0]); + DataTypePtr column_type = make_nullable(argument_types[ElemIdx]); column_data = column_type->create_column(); null_map = &(assert_cast(*column_data).get_null_map_data()); nested_column = assert_cast( @@ -55,12 +56,21 @@ struct AggregateFunctionArrayAggData { } void add(const IColumn& column, size_t row_num) { - const auto& col = assert_cast(column); - const auto& vec = - assert_cast(col.get_nested_column()) - .get_data(); - null_map->push_back(col.get_null_map_data()[row_num]); - nested_column->get_data().push_back(vec[row_num]); + if constexpr (NullableInput) { + const auto& col = assert_cast( + column); + const auto& vec = + assert_cast( + col.get_nested_column()) + .get_data(); + null_map->push_back(col.get_null_map_data()[row_num]); + nested_column->get_data().push_back(vec[row_num]); + } else { + const auto& vec = + assert_cast(column).get_data(); + null_map->push_back(0); + nested_column->get_data().push_back(vec[row_num]); + } DCHECK(null_map->size() == nested_column->size()); } @@ -134,20 +144,20 @@ struct AggregateFunctionArrayAggData { } }; -template +template requires(is_string_type(T)) -struct AggregateFunctionArrayAggData { +struct AggregateFunctionArrayAggData { static constexpr PrimitiveType PType = T; static constexpr bool use_native_serde = false; using ElementType = StringRef; using ColVecType = ColumnString; - using Self = AggregateFunctionArrayAggData; + using Self = AggregateFunctionArrayAggData; MutableColumnPtr column_data; ColVecType* nested_column = nullptr; NullMap* null_map = nullptr; AggregateFunctionArrayAggData(const DataTypes& argument_types) { - DataTypePtr column_type = make_nullable(argument_types[0]); + DataTypePtr column_type = make_nullable(argument_types[ElemIdx]); column_data = column_type->create_column(); null_map = &(assert_cast(*column_data).get_null_map_data()); nested_column = assert_cast( @@ -155,11 +165,19 @@ struct AggregateFunctionArrayAggData { } void add(const IColumn& column, size_t row_num) { - const auto& col = assert_cast(column); - const auto& vec = assert_cast( - col.get_nested_column()); - null_map->push_back(col.get_null_map_data()[row_num]); - nested_column->insert_from(vec, row_num); + if constexpr (NullableInput) { + const auto& col = assert_cast( + column); + const auto& vec = assert_cast( + col.get_nested_column()); + null_map->push_back(col.get_null_map_data()[row_num]); + nested_column->insert_from(vec, row_num); + } else { + const auto& vec = + assert_cast(column); + null_map->push_back(0); + nested_column->insert_from(vec, row_num); + } DCHECK(null_map->size() == nested_column->size()); } @@ -229,18 +247,18 @@ struct AggregateFunctionArrayAggData { } }; -template +template requires(!is_string_type(T) && !is_int_or_bool(T) && !is_float_or_double(T) && !is_decimal(T) && !is_date_type(T) && !is_ip(T)) -struct AggregateFunctionArrayAggData { +struct AggregateFunctionArrayAggData { static constexpr PrimitiveType PType = T; static constexpr bool use_native_serde = true; using ElementType = StringRef; - using Self = AggregateFunctionArrayAggData; + using Self = AggregateFunctionArrayAggData; MutableColumnPtr column_data; AggregateFunctionArrayAggData(const DataTypes& argument_types) - : column_data(argument_types[0]->create_column()) {} + : column_data(argument_types[ElemIdx]->create_column()) {} void add(const IColumn& column, size_t row_num) { column_data->insert_from(column, row_num); } @@ -249,7 +267,16 @@ struct AggregateFunctionArrayAggData { const auto& to_nested_col = to_arr.get_data(); auto start = to_arr.get_offsets()[row_num - 1]; auto end = start + to_arr.get_offsets()[row_num] - to_arr.get_offsets()[row_num - 1]; - column_data->insert_range_from(to_nested_col, start, end - start); + if constexpr (NullableInput) { + column_data->insert_range_from(to_nested_col, start, end - start); + } else { + // Serialized state columns use the nullable result element type while raw-declared + // state stores plain rows; rows serialized this way are all non-null. + const auto& to_nested_nullable = + assert_cast(to_nested_col); + column_data->insert_range_from(to_nested_nullable.get_nested_column(), start, + end - start); + } } void reset() { column_data->clear(); } @@ -257,7 +284,18 @@ struct AggregateFunctionArrayAggData { void insert_result_into(IColumn& to) const { auto& to_arr = assert_cast(to); auto& to_nested_col = to_arr.get_data(); - to_nested_col.insert_range_from(*column_data, 0, column_data->size()); + if constexpr (NullableInput) { + to_nested_col.insert_range_from(*column_data, 0, column_data->size()); + } else { + // The state column stores raw (non-nullable) rows while the result array element + // column is always nullable; copy the raw rows and mark them all non-null. + auto& to_nested_nullable = + assert_cast(to_nested_col); + to_nested_nullable.get_nested_column().insert_range_from( + *column_data, 0, column_data->size()); + to_nested_nullable.get_null_map_data().resize_fill( + to_nested_nullable.get_nested_column().size(), 0); + } to_arr.get_offsets().push_back(to_nested_col.size()); } @@ -482,4 +520,224 @@ class AggregateFunctionArrayAgg final DataTypePtr return_type; }; +/** + * Conditional variant of array_agg: array_agg_if(cond, elem). + * + *

Rows where {@code cond} is false (or NULL) are skipped entirely -- they contribute neither an + * element nor a NULL. Remaining rows follow array_agg semantics exactly (NULL elements are kept). + * Used by IVM to compute insert/delete delta arrays in one aggregate over signed change rows. + */ +template +class AggregateFunctionArrayAggIf final + : public IAggregateFunctionDataHelper, true>, + NotNullableAggregateFunction { +public: + using Base = IAggregateFunctionDataHelper, true>; + + AggregateFunctionArrayAggIf(const DataTypes& argument_types_) + : Base(argument_types_), + return_type( + std::make_shared(make_nullable(argument_types_[1]))) {} + + std::string get_name() const override { return "array_agg_if"; } + + DataTypePtr get_return_type() const override { return return_type; } + + void add(AggregateDataPtr __restrict place, const IColumn** columns, ssize_t row_num, + Arena& arena) const override { + if (!cond_true(columns[0], row_num)) { + return; + } + this->data(place).add(*columns[1], row_num); + } + + void merge(AggregateDataPtr __restrict place, ConstAggregateDataPtr rhs, + Arena& arena) const override { + this->data(place).merge(this->data(rhs)); + } + + void serialize(ConstAggregateDataPtr __restrict place, BufferWritable& buf) const override { + if constexpr (Data::use_native_serde) { + this->data(place).write(buf, *this->argument_types[1], this->version); + } else { + this->data(place).write(buf); + } + } + + void deserialize(AggregateDataPtr __restrict place, BufferReadable& buf, + Arena&) const override { + if constexpr (Data::use_native_serde) { + this->data(place).read(buf, *this->argument_types[1], this->version); + } else { + this->data(place).read(buf); + } + } + + void insert_result_into(ConstAggregateDataPtr __restrict place, IColumn& to) const override { + auto& to_arr = assert_cast(to); + auto& to_nested_col = to_arr.get_data(); + DCHECK(to_nested_col.is_nullable()); + this->data(place).insert_result_into(to); + } + + void serialize_without_key_to_column(ConstAggregateDataPtr __restrict place, + IColumn& to) const override { + this->data(place).insert_result_into(to); + } + + void deserialize_and_merge_vec(const AggregateDataPtr* places, size_t offset, + AggregateDataPtr rhs, const IColumn* column, Arena& arena, + const size_t num_rows) const override { + for (size_t i = 0; i != num_rows; ++i) { + this->data(places[i] + offset).deserialize_and_merge(*column, i); + } + } + + void deserialize_and_merge_from_column_range(AggregateDataPtr __restrict place, + const IColumn& column, size_t begin, size_t end, + Arena& arena) const override { + DCHECK(end <= column.size() && begin <= end) + << ", begin:" << begin << ", end:" << end << ", column.size():" << column.size(); + for (size_t i = begin; i <= end; ++i) { + this->data(place).deserialize_and_merge(column, i); + } + } + + void deserialize_and_merge_vec_selected(const AggregateDataPtr* places, size_t offset, + AggregateDataPtr rhs, const IColumn* column, + Arena& arena, const size_t num_rows) const override { + for (size_t i = 0; i != num_rows; ++i) { + if (places[i]) { + this->data(places[i] + offset).deserialize_and_merge(*column, i); + } + } + } + + void serialize_to_column(const std::vector& places, size_t offset, + MutableColumnPtr& dst, const size_t num_rows) const override { + for (size_t i = 0; i != num_rows; ++i) { + Data& data_ = this->data(places[i] + offset); + data_.insert_result_into(*dst); + } + } + + void streaming_agg_serialize_to_column(const IColumn** columns, MutableColumnPtr& dst, + const size_t num_rows, Arena& arena) const override { + if constexpr (is_string_type(Data::PType)) { + check_array_nullable_string_column_type(*dst, true); + } + auto& to_arr = assert_cast(*dst); + auto& to_nested_col = to_arr.get_data(); + DCHECK(num_rows == columns[0]->size() && num_rows == columns[1]->size()); + auto* col_null = assert_cast(&to_nested_col); + const auto* nullable_elem = + check_and_get_column(*columns[1]); + + for (size_t i = 0; i < num_rows; ++i) { + if (!cond_true(columns[0], i)) { + to_arr.get_offsets().push_back(to_nested_col.size()); + continue; + } + if (nullable_elem != nullptr) { + col_null->get_null_map_data().push_back(nullable_elem->get_null_map_data()[i]); + } else { + col_null->get_null_map_data().push_back(0); + } + const IColumn& elem_src = nullable_elem != nullptr ? nullable_elem->get_nested_column() + : *columns[1]; + if constexpr (is_string_type(Data::PType)) { + auto& vec = assert_cast( + col_null->get_nested_column()); + const auto& vec_src = + assert_cast(elem_src); + vec.insert_from(vec_src, i); + } else if constexpr (!is_string_type(Data::PType) && !is_int_or_bool(Data::PType) && + !is_float_or_double(Data::PType) && !is_decimal(Data::PType) && + !is_date_type(Data::PType) && !is_ip(Data::PType)) { + auto& vec = col_null->get_nested_column(); + vec.insert_from(elem_src, i); + } else { + using ColVecType = typename PrimitiveTypeTraits::ColumnType; + auto& vec = assert_cast( + col_null->get_nested_column()) + .get_data(); + const auto& vec_src = + assert_cast(elem_src) + .get_data(); + vec.push_back(vec_src[i]); + } + to_arr.get_offsets().push_back(to_nested_col.size()); + } + } + + MutableColumnPtr create_serialize_column() const override { + return get_serialized_type()->create_column(); + } + + DataTypePtr get_serialized_type() const override { return return_type; } + + void check_input_columns_type(const IColumn** columns) const override { + IAggregateFunction::check_input_columns_type(columns); + if constexpr (is_string_type(Data::PType)) { + const IColumn& elem_col = *columns[1]; + if (const auto* nullable_column = + check_and_get_column(elem_col)) { + this->template check_argument_column_type( + &nullable_column->get_nested_column()); + } else { + this->template check_argument_column_type(&elem_col); + } + } + } + + void check_result_column_type(const IColumn& to) const override { + IAggregateFunction::check_result_column_type(to); + if constexpr (is_string_type(Data::PType)) { + check_array_nullable_string_column_type(to, true); + } + } + +private: + /** Returns whether the row passes the condition. NULL condition is treated as false. */ + static bool cond_true(const IColumn* cond_column, size_t row_num) { + if (cond_column->is_nullable()) { + const auto& nullable_cond = assert_cast(*cond_column); + if (nullable_cond.is_null_at(row_num)) { + return false; + } + cond_column = &nullable_cond.get_nested_column(); + } + const auto& cond_data = assert_cast( + *cond_column) + .get_data(); + return cond_data[row_num] != 0; + } + + void check_array_nullable_string_column_type(const IColumn& column, + bool is_result_column) const { + const auto* array_column = check_and_get_column(column); + if (UNLIKELY(array_column == nullptr)) { + throw doris::Exception(Status::InternalError( + "Aggregate function {} {} type check failed: Column type {} is not " + "ColumnArray", + get_name(), is_result_column ? "result" : "argument", column.get_name())); + } + + const auto& nested_column = array_column->get_data(); + const auto* nullable_column = check_and_get_column(nested_column); + if (UNLIKELY(nullable_column == nullptr)) { + throw doris::Exception(Status::InternalError( + "Aggregate function {} {} type check failed: Column type {} is not " + "ColumnNullable", + get_name(), is_result_column ? "result" : "argument", + nested_column.get_name())); + } + this->template check_result_column_type_as( + nullable_column->get_nested_column()); + } + + DataTypePtr return_type; +}; + } // namespace doris diff --git a/be/src/exprs/aggregate/aggregate_function_simple_factory.cpp b/be/src/exprs/aggregate/aggregate_function_simple_factory.cpp index 2d1a94dc04c875..aed0980fba4df3 100644 --- a/be/src/exprs/aggregate/aggregate_function_simple_factory.cpp +++ b/be/src/exprs/aggregate/aggregate_function_simple_factory.cpp @@ -64,6 +64,7 @@ void register_aggregate_function_percentile_approx(AggregateFunctionSimpleFactor void register_aggregate_function_orthogonal_bitmap(AggregateFunctionSimpleFactory& factory); void register_aggregate_function_collect_list(AggregateFunctionSimpleFactory& factory); void register_aggregate_function_array_agg(AggregateFunctionSimpleFactory& factory); +void register_aggregate_function_array_agg_if(AggregateFunctionSimpleFactory& factory); void register_aggregate_function_sequence_match(AggregateFunctionSimpleFactory& factory); void register_aggregate_function_avg_weighted(AggregateFunctionSimpleFactory& factory); void register_aggregate_function_histogram(AggregateFunctionSimpleFactory& factory); @@ -119,6 +120,7 @@ AggregateFunctionSimpleFactory& AggregateFunctionSimpleFactory::instance() { register_aggregate_function_orthogonal_bitmap(instance); register_aggregate_function_collect_list(instance); register_aggregate_function_array_agg(instance); + register_aggregate_function_array_agg_if(instance); register_aggregate_function_sequence_match(instance); register_aggregate_function_avg_weighted(instance); register_aggregate_function_histogram(instance); diff --git a/be/test/exprs/aggregate/agg_array_agg_if_test.cpp b/be/test/exprs/aggregate/agg_array_agg_if_test.cpp new file mode 100644 index 00000000000000..c6ed6efad0559a --- /dev/null +++ b/be/test/exprs/aggregate/agg_array_agg_if_test.cpp @@ -0,0 +1,227 @@ +// 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 +#include + +#include +#include +#include +#include + +#include "core/column/column.h" +#include "core/column/column_array.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 "core/types.h" +#include "exprs/aggregate/agg_function_test.h" +#include "gtest/gtest_pred_impl.h" + +namespace doris { + +struct AggregateFunctionArrayAggIfTest : public AggregateFunctiontest {}; + +namespace { + +DataTypePtr bool_type() { + return std::make_shared(); +} + +DataTypePtr nullable_bool_type() { + return make_nullable(std::make_shared()); +} + +DataTypePtr nullable_int64_type() { + return make_nullable(std::make_shared()); +} + +DataTypePtr nullable_string_type() { + return make_nullable(std::make_shared()); +} + +MutableColumnPtr make_cond_column(std::initializer_list conds) { + auto cond = ColumnUInt8::create(); + for (int c : conds) { + cond->insert_value(c != 0); + } + return cond; +} + +MutableColumnPtr make_nullable_cond_column(std::initializer_list conds) { + auto values = ColumnUInt8::create(); + auto null_map = ColumnUInt8::create(); + for (const char* c : conds) { + if (*c == 'n') { + values->insert_value(0); + null_map->insert_value(1); + } else { + values->insert_value(*c == 't'); + null_map->insert_value(0); + } + } + return ColumnNullable::create(std::move(values), std::move(null_map)); +} + +/** Runs array_agg_if over the given cond/elem block and checks the aggregated array row. */ +void check_array_agg_if(AggregateFunctionArrayAggIfTest* test, Block block, + DataTypePtr cond_type, DataTypePtr elem_type, Array expected) { + test->create_agg("array_agg_if", false, {cond_type, elem_type}, elem_type); + auto array_type = std::make_shared(make_nullable(elem_type)); + auto expected_column = array_type->create_column(); + expected_column->insert(Field::create_field(std::move(expected))); + test->execute(std::move(block), + ColumnWithTypeAndName(std::move(expected_column), array_type, "column")); +} + +Block make_block(MutableColumnPtr cond, MutableColumnPtr elem, DataTypePtr cond_type, + DataTypePtr elem_type) { + return Block({ColumnWithTypeAndName(std::move(cond), cond_type, "cond"), + ColumnWithTypeAndName(std::move(elem), elem_type, "elem")}); +} + +} // namespace + +TEST_F(AggregateFunctionArrayAggIfTest, test_int64_skip_cond_false_rows) { + auto elem_type = nullable_int64_type(); + auto elem = elem_type->create_column(); + elem->insert(Field::create_field(1)); + elem->insert(Field::create_field(2)); + elem->insert(Field::create_field(3)); + elem->insert(Field::create_field(4)); + auto block = make_block(make_cond_column({1, 0, 1, 1}), std::move(elem), bool_type(), + elem_type); + check_array_agg_if(this, std::move(block), bool_type(), elem_type, + Array {Field::create_field(1), + Field::create_field(3), + Field::create_field(4)}); +} + +TEST_F(AggregateFunctionArrayAggIfTest, test_int64_keeps_null_elements) { + auto elem_type = nullable_int64_type(); + auto elem = elem_type->create_column(); + elem->insert(Field::create_field(5)); + elem->insert(Field()); + elem->insert(Field::create_field(7)); + auto block = make_block(make_cond_column({1, 1, 1}), std::move(elem), bool_type(), + elem_type); + check_array_agg_if(this, std::move(block), bool_type(), elem_type, + Array {Field::create_field(5), Field(), + Field::create_field(7)}); +} + +TEST_F(AggregateFunctionArrayAggIfTest, test_int64_null_elem_skipped_by_cond) { + auto elem_type = nullable_int64_type(); + auto elem = elem_type->create_column(); + elem->insert(Field::create_field(5)); + elem->insert(Field()); + elem->insert(Field::create_field(7)); + auto block = make_block(make_cond_column({1, 0, 1}), std::move(elem), bool_type(), + elem_type); + check_array_agg_if(this, std::move(block), bool_type(), elem_type, + Array {Field::create_field(5), + Field::create_field(7)}); +} + +TEST_F(AggregateFunctionArrayAggIfTest, test_int64_all_cond_false_returns_empty_array) { + auto elem_type = nullable_int64_type(); + auto elem = elem_type->create_column(); + elem->insert(Field::create_field(1)); + elem->insert(Field::create_field(2)); + auto block = make_block(make_cond_column({0, 0}), std::move(elem), bool_type(), elem_type); + check_array_agg_if(this, std::move(block), bool_type(), elem_type, Array {}); +} + +TEST_F(AggregateFunctionArrayAggIfTest, test_int64_null_cond_treated_as_false) { + auto elem_type = nullable_int64_type(); + auto elem = elem_type->create_column(); + elem->insert(Field::create_field(1)); + elem->insert(Field::create_field(2)); + elem->insert(Field::create_field(3)); + auto block = make_block(make_nullable_cond_column({"t", "n", "f"}), std::move(elem), + nullable_bool_type(), elem_type); + check_array_agg_if(this, std::move(block), nullable_bool_type(), elem_type, + Array {Field::create_field(1)}); +} + +TEST_F(AggregateFunctionArrayAggIfTest, test_string_keeps_null_elements) { + auto elem_type = nullable_string_type(); + auto elem = elem_type->create_column(); + elem->insert(Field::create_field(String("a"))); + elem->insert(Field()); + elem->insert(Field::create_field(String("c"))); + auto block = make_block(make_cond_column({1, 1, 1}), std::move(elem), bool_type(), + elem_type); + check_array_agg_if(this, std::move(block), bool_type(), elem_type, + Array {Field::create_field(String("a")), Field(), + Field::create_field(String("c"))}); +} + +namespace { + +Field int_array_field(std::initializer_list values) { + Array fields; + for (int v : values) { + fields.push_back(Field::create_field(v)); + } + return Field::create_field(std::move(fields)); +} + +} // namespace + +TEST_F(AggregateFunctionArrayAggIfTest, test_complex_elem_not_nullable) { + // Outer element type is a non-nullable ARRAY: exercises the raw native-serde state path. + auto inner_type = make_nullable(std::make_shared()); + auto elem_type = std::make_shared(inner_type); + auto elem = elem_type->create_column(); + elem->insert(int_array_field({1, 2})); + elem->insert(int_array_field({3})); + auto block = make_block(make_cond_column({1, 0}), std::move(elem), bool_type(), elem_type); + check_array_agg_if(this, std::move(block), bool_type(), elem_type, + Array {int_array_field({1, 2})}); +} + +TEST_F(AggregateFunctionArrayAggIfTest, test_complex_elem_nullable) { + auto inner_type = make_nullable(std::make_shared()); + auto elem_type = make_nullable(std::make_shared(inner_type)); + auto elem = elem_type->create_column(); + elem->insert(int_array_field({1, 2})); + elem->insert(Field()); + elem->insert(int_array_field({3})); + auto block = make_block(make_cond_column({1, 0, 1}), std::move(elem), bool_type(), elem_type); + check_array_agg_if(this, std::move(block), bool_type(), elem_type, + Array {int_array_field({1, 2}), int_array_field({3})}); +} + +TEST_F(AggregateFunctionArrayAggIfTest, test_string_skip_cond_false_rows) { + auto elem_type = nullable_string_type(); + auto elem = elem_type->create_column(); + elem->insert(Field::create_field(String("a"))); + elem->insert(Field::create_field(String("b"))); + elem->insert(Field::create_field(String("c"))); + auto block = make_block(make_cond_column({1, 0, 1}), std::move(elem), bool_type(), + elem_type); + check_array_agg_if(this, std::move(block), bool_type(), elem_type, + Array {Field::create_field(String("a")), + Field::create_field(String("c"))}); +} + +} // namespace doris diff --git a/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinAggregateFunctions.java b/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinAggregateFunctions.java index 7f3fa4cddec3bc..8d9b5b0902ded9 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinAggregateFunctions.java +++ b/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinAggregateFunctions.java @@ -20,6 +20,7 @@ import org.apache.doris.nereids.trees.expressions.functions.agg.AIAgg; import org.apache.doris.nereids.trees.expressions.functions.agg.AnyValue; import org.apache.doris.nereids.trees.expressions.functions.agg.ArrayAgg; +import org.apache.doris.nereids.trees.expressions.functions.agg.ArrayAggIf; import org.apache.doris.nereids.trees.expressions.functions.agg.Avg; import org.apache.doris.nereids.trees.expressions.functions.agg.AvgMap; import org.apache.doris.nereids.trees.expressions.functions.agg.AvgWeighted; @@ -133,6 +134,7 @@ private BuiltinAggregateFunctions() { agg(AIAgg.class, "ai_agg"), agg(AnyValue.class, "any", "any_value"), agg(ArrayAgg.class, "array_agg"), + agg(ArrayAggIf.class, "array_agg_if"), agg(Avg.class, "avg"), agg(AvgMap.class, "avg_map"), agg(AvgWeighted.class, "avg_weighted"), diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/ArrayAggIf.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/ArrayAggIf.java new file mode 100644 index 00000000000000..02cb14c0579bab --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/ArrayAggIf.java @@ -0,0 +1,78 @@ +// 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. + +package org.apache.doris.nereids.trees.expressions.functions.agg; + +import org.apache.doris.catalog.FunctionSignature; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature; +import org.apache.doris.nereids.trees.expressions.literal.ArrayLiteral; +import org.apache.doris.nereids.trees.expressions.shape.BinaryExpression; +import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor; +import org.apache.doris.nereids.types.ArrayType; +import org.apache.doris.nereids.types.BooleanType; +import org.apache.doris.nereids.types.coercion.AnyDataType; +import org.apache.doris.nereids.types.coercion.FollowToAnyDataType; + +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; + +import java.util.ArrayList; +import java.util.List; + +/** + * AggregateFunction 'array_agg_if'. Collects elements of rows whose condition is true, + * following array_agg semantics (NULL elements are kept); rows whose condition is false + * (or NULL) are skipped entirely. + */ +public class ArrayAggIf extends NotNullableAggregateFunction + implements BinaryExpression, ExplicitlyCastableSignature { + public static final List SIGNATURES = ImmutableList.of( + FunctionSignature.ret(ArrayType.of(new FollowToAnyDataType(1))) + .args(BooleanType.INSTANCE, new AnyDataType(1)) + ); + + public ArrayAggIf(Expression cond, Expression elem) { + super("array_agg_if", cond, elem); + } + + /** constructor for withChildren and reuse signature */ + private ArrayAggIf(AggregateFunctionParams functionParams) { + super(functionParams); + } + + @Override + public AggregateFunction withDistinctAndChildren(boolean distinct, List children) { + Preconditions.checkArgument(children.size() == 2); + return new ArrayAggIf(getFunctionParams(distinct, children)); + } + + @Override + public R accept(ExpressionVisitor visitor, C context) { + return visitor.visitArrayAggIf(this, context); + } + + @Override + public List getSignatures() { + return SIGNATURES; + } + + @Override + public Expression resultForEmptyInput() { + return new ArrayLiteral(new ArrayList<>(), this.getDataType()); + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/AggregateFunctionVisitor.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/AggregateFunctionVisitor.java index 8e0151b6bdf7e2..4fe49aea0eb9e6 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/AggregateFunctionVisitor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/AggregateFunctionVisitor.java @@ -21,6 +21,7 @@ import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction; import org.apache.doris.nereids.trees.expressions.functions.agg.AnyValue; import org.apache.doris.nereids.trees.expressions.functions.agg.ArrayAgg; +import org.apache.doris.nereids.trees.expressions.functions.agg.ArrayAggIf; import org.apache.doris.nereids.trees.expressions.functions.agg.Avg; import org.apache.doris.nereids.trees.expressions.functions.agg.AvgWeighted; import org.apache.doris.nereids.trees.expressions.functions.agg.BitmapAgg; @@ -125,6 +126,10 @@ default R visitArrayAgg(ArrayAgg arrayAgg, C context) { return visitAggregateFunction(arrayAgg, context); } + default R visitArrayAggIf(ArrayAggIf arrayAggIf, C context) { + return visitAggregateFunction(arrayAggIf, context); + } + default R visitAvg(Avg avg, C context) { return visitNullableAggregateFunction(avg, context); } diff --git a/regression-test/data/nereids_function_p0/agg_function/agg.out b/regression-test/data/nereids_function_p0/agg_function/agg.out index d0f3dfc22a60d3..2be041b1e8eef8 100644 --- a/regression-test/data/nereids_function_p0/agg_function/agg.out +++ b/regression-test/data/nereids_function_p0/agg_function/agg.out @@ -6996,3 +6996,46 @@ true 10 [11] 11 [12] +-- !sql_array_agg_if -- +\N [] +0 [] +1 [] +2 [] +3 [] +4 [] +5 [] +6 [] +7 [8] +8 [9] +9 [10] +10 [11] +11 [12] + +-- !sql_array_agg_if_not_nullable -- +0 [] +1 [] +2 [] +3 [] +4 [] +5 [] +6 [] +7 [8] +8 [9] +9 [10] +10 [11] +11 [12] + +-- !sql_array_agg_if_expr_cond -- +\N [] +0 [1] +1 [2] +2 [3] +3 [4] +4 [5] +5 [6] +6 [7] +7 [8] +8 [9] +9 [10] +10 [11] +11 [12] diff --git a/regression-test/suites/nereids_function_p0/agg_function/agg.groovy b/regression-test/suites/nereids_function_p0/agg_function/agg.groovy index 6ac820aa47a917..c76caafaaff810 100644 --- a/regression-test/suites/nereids_function_p0/agg_function/agg.groovy +++ b/regression-test/suites/nereids_function_p0/agg_function/agg.groovy @@ -3214,4 +3214,13 @@ suite("nereids_agg_fn") { qt_sql_array_agg_not_nullable ''' select id,array_agg(kint) from fn_test_not_nullable group by id order by id''' + qt_sql_array_agg_if ''' + select id,array_agg_if(kbool,kint) from fn_test group by id order by id''' + + qt_sql_array_agg_if_not_nullable ''' + select id,array_agg_if(kbool,kint) from fn_test_not_nullable group by id order by id''' + + qt_sql_array_agg_if_expr_cond ''' + select id,array_agg_if(kint > 0,kint) from fn_test group by id order by id''' + }