diff --git a/be/src/exprs/aggregate/aggregate_function_distinct.h b/be/src/exprs/aggregate/aggregate_function_distinct.h index c1dfb6bae5d0ce..c2c84934c30c58 100644 --- a/be/src/exprs/aggregate/aggregate_function_distinct.h +++ b/be/src/exprs/aggregate/aggregate_function_distinct.h @@ -72,7 +72,13 @@ struct AggregateFunctionDistinctSingleNumericData { void merge(const Self& rhs, Arena&) { DCHECK(!stable); if constexpr (!stable) { - data.merge(Container(rhs.data)); + // Only an empty destination has a known final size; other sets may overlap. + if (data.empty() && !rhs.data.empty()) { + data.reserve(rhs.data.size()); + } + for (const auto& elem : rhs.data) { + data.insert(elem); + } } } @@ -91,6 +97,10 @@ struct AggregateFunctionDistinctSingleNumericData { if constexpr (!stable) { uint64_t new_size = 0; buf.read_var_uint(new_size); + // Avoid reserving an upper bound when merging into a nonempty set. + if (data.empty() && new_size != 0) { + data.reserve(new_size); + } typename PrimitiveTypeTraits::CppType x; for (size_t i = 0; i < new_size; ++i) { buf.read_binary(x); @@ -99,6 +109,11 @@ struct AggregateFunctionDistinctSingleNumericData { } } + void deserialize_and_merge(Self& /*rhs*/, BufferReadable& buf, Arena& arena) { + // Numeric keys can be inserted directly without building a temporary hash set. + deserialize(buf, arena); + } + MutableColumns get_arguments(const DataTypes& argument_types) const { MutableColumns argument_columns; argument_columns.emplace_back(argument_types[0]->create_column()); @@ -153,6 +168,12 @@ struct AggregateFunctionDistinctGenericData { } } + void deserialize_and_merge(Self& rhs, BufferReadable& buf, Arena& arena) { + // deserialize() borrows StringRefs from buf; merge() copies them into the arena. + rhs.deserialize(buf, arena); + merge(rhs, arena); + } + void deserialize(BufferReadable& buf, Arena& arena) { DCHECK(!stable); if constexpr (!stable) { @@ -308,6 +329,11 @@ class AggregateFunctionDistinct final this->data(place).deserialize(buf, arena); } + void deserialize_and_merge(AggregateDataPtr __restrict place, AggregateDataPtr __restrict rhs, + BufferReadable& buf, Arena& arena) const override { + this->data(place).deserialize_and_merge(this->data(rhs), buf, arena); + } + void insert_result_into(ConstAggregateDataPtr targetplace, IColumn& to) const override { // place is essentially an AggregateDataPtr, passed as a ConstAggregateDataPtr. auto* place = const_cast(targetplace); diff --git a/be/test/exprs/aggregate/aggregate_function_distinct_test.cpp b/be/test/exprs/aggregate/aggregate_function_distinct_test.cpp new file mode 100644 index 00000000000000..132c332990bb46 --- /dev/null +++ b/be/test/exprs/aggregate/aggregate_function_distinct_test.cpp @@ -0,0 +1,247 @@ +// 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 "exprs/aggregate/aggregate_function_distinct.h" + +#include + +#include +#include +#include +#include +#include + +#include "agent/be_exec_version_manager.h" +#include "core/arena.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_nullable.h" +#include "core/data_type/data_type_number.h" +#include "core/data_type/data_type_string.h" +#include "core/string_buffer.hpp" +#include "exec/common/hash_table/phmap_fwd_decl.h" +#include "exprs/aggregate/aggregate_function_simple_factory.h" +#include "testutil/column_helper.h" + +namespace doris { +namespace { + +template +class DistinctNumericMergeTest : public testing::Test {}; + +using IntegerTypes = testing::Types, + std::integral_constant, + std::integral_constant, + std::integral_constant, + std::integral_constant>; +TYPED_TEST_SUITE(DistinctNumericMergeTest, IntegerTypes); + +TYPED_TEST(DistinctNumericMergeTest, PreserveSourceAndDeduplicateRepeatedMerges) { + using Data = AggregateFunctionDistinctSingleNumericData; + Arena arena; + Data destination; + Data source; + source.data.insert({1, 2, 3}); + + destination.merge(source, arena); + EXPECT_EQ(destination.data, source.data); + destination.data.insert(4); + const auto capacity = destination.data.capacity(); + destination.merge(source, arena); + EXPECT_EQ(destination.data.size(), 4); + EXPECT_EQ(destination.data.capacity(), capacity); + EXPECT_EQ(source.data.size(), 3); + EXPECT_FALSE(source.data.contains(4)); + + Data empty; + destination.merge(empty, arena); + EXPECT_EQ(destination.data.size(), 4); + EXPECT_EQ(destination.data.capacity(), capacity); + destination.clear(); + destination.merge(empty, arena); + EXPECT_TRUE(destination.data.empty()); +} + +TYPED_TEST(DistinctNumericMergeTest, DeserializeIntoExistingDestinationWithoutPopulatingScratch) { + using Data = AggregateFunctionDistinctSingleNumericData; + Arena arena; + Data source; + source.data.insert({1, 2, 3}); + ColumnString serialized; + VectorBufferWriter writer(serialized); + source.serialize(writer); + writer.commit(); + + Data destination; + Data scratch; + destination.data.insert(4); + for (int repeat = 0; repeat < 2; ++repeat) { + VectorBufferReader reader(serialized.get_data_at(0)); + destination.deserialize_and_merge(scratch, reader, arena); + EXPECT_EQ(destination.data.size(), 4); + for (int value : {1, 2, 3, 4}) { + EXPECT_TRUE(destination.data.contains(value)); + } + // The optimization must not materialize the serialized set in the scratch state. + EXPECT_TRUE(scratch.data.empty()); + EXPECT_EQ(scratch.data.capacity(), 0); + } + + Data empty; + ColumnString serialized_empty; + VectorBufferWriter empty_writer(serialized_empty); + empty.serialize(empty_writer); + empty_writer.commit(); + VectorBufferReader empty_reader(serialized_empty.get_data_at(0)); + destination.deserialize_and_merge(scratch, empty_reader, arena); + EXPECT_EQ(destination.data.size(), 4); + + destination.clear(); + VectorBufferReader reader(serialized.get_data_at(0)); + destination.deserialize_and_merge(scratch, reader, arena); + EXPECT_EQ(destination.data, source.data); + EXPECT_TRUE(scratch.data.empty()); +} + +void add_column(const IAggregateFunction& function, AggregateDataPtr place, const IColumn& column, + Arena& arena) { + const IColumn* columns[] = {&column}; + function.add_batch_single_place(column.size(), place, columns, arena); +} + +class DistinctMergeDispatchTest : public testing::TestWithParam { +protected: + AggregateFunctionPtr function(const std::string& name, const DataTypePtr& input_type, + const DataTypePtr& result_type, bool result_nullable) { + AggregateFunctionAttr attr; + attr.enable_aggregate_function_null_v2 = GetParam(); + return AggregateFunctionSimpleFactory::instance().get( + name, {input_type}, result_type, result_nullable, + BeExecVersionManager::get_newest_version(), attr); + } +}; + +TEST_P(DistinctMergeDispatchTest, NullableSumMergesColumnRangesAndAllNullStates) { + auto type = make_nullable(std::make_shared()); + auto aggregate = function("multi_distinct_sum", type, type, true); + ASSERT_NE(aggregate, nullptr); + Arena arena; + AggregateFunctionGuard source(aggregate.get()); + AggregateFunctionGuard destination(aggregate.get()); + auto serialized = aggregate->get_serialized_type()->create_column(); + + auto first = ColumnHelper::create_nullable_column({1, 2, 2, 99}, {0, 0, 0, 1}); + add_column(*aggregate, source.data(), *first, arena); + aggregate->serialize_without_key_to_column(source.data(), *serialized); + aggregate->reset(source.data()); + auto second = ColumnHelper::create_nullable_column({2, 3, 99}, {0, 0, 1}); + add_column(*aggregate, source.data(), *second, arena); + aggregate->serialize_without_key_to_column(source.data(), *serialized); + aggregate->reset(source.data()); + auto nulls = ColumnHelper::create_nullable_column({99, 99}, {1, 1}); + add_column(*aggregate, source.data(), *nulls, arena); + aggregate->serialize_without_key_to_column(source.data(), *serialized); + + auto initial = ColumnHelper::create_nullable_column({10}, {0}); + add_column(*aggregate, destination.data(), *initial, arena); + for (int repeat = 0; repeat < 2; ++repeat) { + aggregate->deserialize_and_merge_from_column(destination.data(), *serialized, arena); + } + auto result = type->create_column(); + aggregate->insert_result_into(destination.data(), *result); + EXPECT_TRUE(ColumnHelper::column_equal( + std::move(result), ColumnHelper::create_nullable_column({16}, {0}))); + + aggregate->reset(destination.data()); + aggregate->deserialize_and_merge_from_column_range(destination.data(), *serialized, 2, 2, + arena); + result = type->create_column(); + aggregate->insert_result_into(destination.data(), *result); + EXPECT_TRUE(ColumnHelper::column_equal( + std::move(result), ColumnHelper::create_nullable_column({0}, {1}))); +} + +TEST_P(DistinctMergeDispatchTest, GroupedAndSelectedBatchMerges) { + for (bool nullable : {false, true}) { + DataTypePtr type = std::make_shared(); + if (nullable) { + type = make_nullable(type); + } + auto aggregate = function("multi_distinct_sum", type, type, nullable); + ASSERT_NE(aggregate, nullptr); + Arena arena; + AggregateFunctionGuard source(aggregate.get()); + AggregateFunctionGuard destination(aggregate.get()); + auto serialized = aggregate->get_serialized_type()->create_column(); + for (const auto& values : {std::vector {1, 2, 2}, std::vector {2, 3}}) { + auto column = type->create_column(); + for (auto value : values) { + column->insert(Field::create_field(value)); + } + aggregate->reset(source.data()); + add_column(*aggregate, source.data(), *column, arena); + aggregate->serialize_without_key_to_column(source.data(), *serialized); + } + + auto* scratch = reinterpret_cast( + arena.aligned_alloc(2 * aggregate->size_of_data(), aggregate->align_of_data())); + std::array places {destination.data(), destination.data()}; + aggregate->deserialize_and_merge_vec(places.data(), 0, scratch, serialized.get(), arena, 2); + auto result = type->create_column(); + aggregate->insert_result_into(destination.data(), *result); + EXPECT_EQ(type->to_string(*result, 0), "6"); + + aggregate->reset(destination.data()); + places[0] = nullptr; + aggregate->deserialize_and_merge_vec_selected(places.data(), 0, scratch, serialized.get(), + arena, 2); + result = type->create_column(); + aggregate->insert_result_into(destination.data(), *result); + EXPECT_EQ(type->to_string(*result, 0), "5"); + } +} + +TEST_P(DistinctMergeDispatchTest, GenericStringStateOwnsKeysAfterInputBufferIsReleased) { + auto type = std::make_shared(); + auto aggregate = function("multi_distinct_min", type, type, false); + ASSERT_NE(aggregate, nullptr); + auto serialized = ColumnString::create(); + const std::string first(512, 'a'); + const std::string second(512, 'b'); + { + Arena source_arena; + AggregateFunctionGuard source(aggregate.get()); + auto column = ColumnHelper::create_column({first, second, first}); + add_column(*aggregate, source.data(), *column, source_arena); + aggregate->serialize_without_key_to_column(source.data(), *serialized); + } + + Arena arena; + AggregateFunctionGuard destination(aggregate.get()); + aggregate->deserialize_and_merge_from_column(destination.data(), *serialized, arena); + std::fill(serialized->get_chars().begin(), serialized->get_chars().end(), '?'); + serialized.reset(); + auto result = type->create_column(); + aggregate->insert_result_into(destination.data(), *result); + EXPECT_EQ(result->get_data_at(0).to_string(), first); +} + +INSTANTIATE_TEST_SUITE_P(NullImplementations, DistinctMergeDispatchTest, testing::Bool()); + +} // namespace +} // namespace doris diff --git a/regression-test/data/query_p0/aggregate/test_numeric_distinct_merge.out b/regression-test/data/query_p0/aggregate/test_numeric_distinct_merge.out new file mode 100644 index 00000000000000..bbfe6f8e6c7e8d --- /dev/null +++ b/regression-test/data/query_p0/aggregate/test_numeric_distinct_merge.out @@ -0,0 +1,37 @@ +-- This file is automatically generated. You should know what you did if you want to edit this +-- !integer_types_1 -- +9 9 9 9 9 9 + +-- !grouped_1 -- +\N 4 4 +0 3 3 +1 4 4 +2 \N 0 + +-- !all_null_1 -- +\N + +-- !empty_1 -- +\N \N + +-- !state_merge_1 -- +9 + +-- !integer_types_2 -- +9 9 9 9 9 9 + +-- !grouped_2 -- +\N 4 4 +0 3 3 +1 4 4 +2 \N 0 + +-- !all_null_2 -- +\N + +-- !empty_2 -- +\N \N + +-- !state_merge_2 -- +9 + diff --git a/regression-test/suites/query_p0/aggregate/test_numeric_distinct_merge.groovy b/regression-test/suites/query_p0/aggregate/test_numeric_distinct_merge.groovy new file mode 100644 index 00000000000000..5c6b765dfe372e --- /dev/null +++ b/regression-test/suites/query_p0/aggregate/test_numeric_distinct_merge.groovy @@ -0,0 +1,70 @@ +// 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. + +suite("test_numeric_distinct_merge") { + sql "DROP TABLE IF EXISTS test_numeric_distinct_merge" + sql """ + CREATE TABLE test_numeric_distinct_merge ( + id INT NOT NULL, + g INT NULL, + v BIGINT NULL, + nonnull_v BIGINT NOT NULL + ) DUPLICATE KEY(id) + DISTRIBUTED BY HASH(id) BUCKETS 4 + PROPERTIES ("replication_num" = "1") + """ + sql """ + INSERT INTO test_numeric_distinct_merge VALUES + (1, 0, 1, 1), (2, 0, 2, 2), (3, 0, 2, 2), (4, 0, NULL, 0), + (5, 1, 2, 2), (6, 1, 3, 3), (7, 1, -1, -1), (8, 1, NULL, 0), + (9, 2, NULL, 0), (10, NULL, 4, 4), (11, NULL, 4, 4), (12, NULL, NULL, 0) + """ + + for (phase in [1, 2]) { + sql "SET agg_phase = ${phase}" + + "order_qt_integer_types_${phase}" """ + SELECT multi_distinct_sum(CAST(v AS TINYINT)), + multi_distinct_sum(CAST(v AS SMALLINT)), + multi_distinct_sum(CAST(v AS INT)), + multi_distinct_sum(v), + multi_distinct_sum(CAST(v AS LARGEINT)), + multi_distinct_sum(nonnull_v) + FROM test_numeric_distinct_merge + """ + "order_qt_grouped_${phase}" """ + SELECT g, multi_distinct_sum(v), multi_distinct_sum(nonnull_v) + FROM test_numeric_distinct_merge GROUP BY g + """ + "order_qt_all_null_${phase}" """ + SELECT multi_distinct_sum(v) FROM test_numeric_distinct_merge WHERE g = 2 + """ + "order_qt_empty_${phase}" """ + SELECT multi_distinct_sum(v), multi_distinct_sum(nonnull_v) + FROM test_numeric_distinct_merge WHERE id < 0 + """ + "order_qt_state_merge_${phase}" """ + SELECT multi_distinct_sum_merge(multi_distinct_sum_state(v)) + FROM test_numeric_distinct_merge + """ + } + + explain { + sql "SELECT multi_distinct_sum(v) FROM test_numeric_distinct_merge" + contains "merge finalize" + } +}