github-actions[bot] commented on code in PR #69456: URL: https://github.com/apache/doris/pull/69456#discussion_r4400070091
########## be/src/exprs/aggregate/aggregate_function_topk.h: ########## @@ -0,0 +1,741 @@ +// 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. + +#pragma once + +#include <algorithm> +#include <cstdint> +#include <limits> +#include <memory> +#include <type_traits> +#include <utility> + +#include "common/check.h" +#include "common/compare.h" +#include "common/exception.h" +#include "core/assert_cast.h" +#include "core/column/column.h" +#include "core/column/column_array.h" +#include "core/column/column_decimal.h" +#include "core/column/column_nullable.h" +#include "core/column/column_string.h" +#include "core/column/column_struct.h" +#include "core/column/column_vector.h" +#include "core/custom_allocator.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 "core/data_type/data_type_struct.h" +#include "core/data_type/primitive_type.h" +#include "core/memcmp_small.h" +#include "core/pod_array.h" +#include "core/string_buffer.hpp" +#include "core/string_ref.h" +#include "core/types.h" +#include "exprs/aggregate/aggregate_function.h" + +namespace doris { + +/// The values of one group, in arrival order, one entry per slot. +/// +/// They deliberately do not live in a Doris column: a column's container allocates a 4 KiB initial +/// chunk (PaddedPODArray) as soon as it holds one value, so one column per group costs ~4 KiB per +/// group regardless of k. That dominates the memory of a GROUP BY with millions of groups, while +/// these stores grow with the number of values actually retained. +template <typename T> +struct FixedTopKValueStore { + DorisVector<T> values; + + size_t size() const { return values.size(); } + void reserve(size_t n) { values.reserve(n); } + void clear() { values.clear(); } + const T& at(size_t slot) const { return values[slot]; } + void append(const T& value) { values.push_back(value); } + void swap(FixedTopKValueStore& other) noexcept { values.swap(other.values); } +}; + +/// A StringRef cannot be stored: the arena reallocates and would invalidate the pointers. A slot +/// keeps an offset/length into the bytes instead and materializes a view on access. +struct StringTopKValueStore { + struct Span { + uint32_t offset = 0; + uint32_t length = 0; + }; + + /// The initial allocation is small on purpose; see the comment on FixedTopKValueStore. + PODArray<char, 32> bytes; Review Comment: [P1] Pad the string store before overflow-reading comparisons For `min_topk(s, 1)` over two rows with `s='a'`, the first append allocates a one-byte buffer: `PODArray::insert` calls `reserve` directly, so `initial_bytes=32` provides no padding. The next add compares the stored `StringRef` through `memcmp_small_allow_overflow15`, whose SSE2 branch reads 16 bytes even for a one-byte string. That load extends 15 bytes beyond the allocation, violating the comparator's memory contract and ASAN bounds. `max_topk`, `k=0` and state merges share this path. Use storage with at least 15 readable bytes after the payload, or a bounds-safe comparison, and cover short strings in aggregate/state-merge tests. This is established from the allocation and load paths; no ASAN run was performed. ########## be/src/exprs/aggregate/aggregate_function_topk.h: ########## @@ -0,0 +1,741 @@ +// 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. + +#pragma once + +#include <algorithm> +#include <cstdint> +#include <limits> +#include <memory> +#include <type_traits> +#include <utility> + +#include "common/check.h" +#include "common/compare.h" +#include "common/exception.h" +#include "core/assert_cast.h" +#include "core/column/column.h" +#include "core/column/column_array.h" +#include "core/column/column_decimal.h" +#include "core/column/column_nullable.h" +#include "core/column/column_string.h" +#include "core/column/column_struct.h" +#include "core/column/column_vector.h" +#include "core/custom_allocator.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 "core/data_type/data_type_struct.h" +#include "core/data_type/primitive_type.h" +#include "core/memcmp_small.h" +#include "core/pod_array.h" +#include "core/string_buffer.hpp" +#include "core/string_ref.h" +#include "core/types.h" +#include "exprs/aggregate/aggregate_function.h" + +namespace doris { + +/// The values of one group, in arrival order, one entry per slot. +/// +/// They deliberately do not live in a Doris column: a column's container allocates a 4 KiB initial +/// chunk (PaddedPODArray) as soon as it holds one value, so one column per group costs ~4 KiB per +/// group regardless of k. That dominates the memory of a GROUP BY with millions of groups, while +/// these stores grow with the number of values actually retained. +template <typename T> +struct FixedTopKValueStore { + DorisVector<T> values; + + size_t size() const { return values.size(); } + void reserve(size_t n) { values.reserve(n); } + void clear() { values.clear(); } + const T& at(size_t slot) const { return values[slot]; } + void append(const T& value) { values.push_back(value); } + void swap(FixedTopKValueStore& other) noexcept { values.swap(other.values); } +}; + +/// A StringRef cannot be stored: the arena reallocates and would invalidate the pointers. A slot +/// keeps an offset/length into the bytes instead and materializes a view on access. +struct StringTopKValueStore { + struct Span { + uint32_t offset = 0; + uint32_t length = 0; + }; + + /// The initial allocation is small on purpose; see the comment on FixedTopKValueStore. + PODArray<char, 32> bytes; + DorisVector<Span> spans; + + size_t size() const { return spans.size(); } + void reserve(size_t n) { spans.reserve(n); } + void clear() { + bytes.clear(); + spans.clear(); + } + StringRef at(size_t slot) const { + const Span& span = spans[slot]; + return StringRef(bytes.data() + span.offset, span.length); + } + void append(StringRef value) { + DCHECK_LE(bytes.size() + value.size, std::numeric_limits<uint32_t>::max()); + DCHECK_LE(value.size, std::numeric_limits<uint32_t>::max()); + spans.push_back({static_cast<uint32_t>(bytes.size()), static_cast<uint32_t>(value.size)}); Review Comment: [P2] Prevent cumulative string offsets from wrapping For `k=2049`, one group receiving 4097 descending distinct 1 MiB strings across input blocks reaches this append with `bytes.size()==2^32`. Only 2047 slots have been evicted before this append, so `dead > live` has never triggered compaction. With enough query memory, release builds narrow the new offset to zero: the newly retained smallest value then reads the first, already-evicted high string. Checked builds instead fail the size DCHECK at the preceding append. The live result is only 2049 MiB, so the per-column 4 GiB limit does not exclude this schedule. Use lossless offsets or compact/check the byte limit before narrowing, returning a query error for unsupported sizes. ########## be/src/exprs/aggregate/aggregate_function_topk.h: ########## @@ -0,0 +1,741 @@ +// 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. + +#pragma once + +#include <algorithm> +#include <cstdint> +#include <limits> +#include <memory> +#include <type_traits> +#include <utility> + +#include "common/check.h" +#include "common/compare.h" +#include "common/exception.h" +#include "core/assert_cast.h" +#include "core/column/column.h" +#include "core/column/column_array.h" +#include "core/column/column_decimal.h" +#include "core/column/column_nullable.h" +#include "core/column/column_string.h" +#include "core/column/column_struct.h" +#include "core/column/column_vector.h" +#include "core/custom_allocator.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 "core/data_type/data_type_struct.h" +#include "core/data_type/primitive_type.h" +#include "core/memcmp_small.h" +#include "core/pod_array.h" +#include "core/string_buffer.hpp" +#include "core/string_ref.h" +#include "core/types.h" +#include "exprs/aggregate/aggregate_function.h" + +namespace doris { + +/// The values of one group, in arrival order, one entry per slot. +/// +/// They deliberately do not live in a Doris column: a column's container allocates a 4 KiB initial +/// chunk (PaddedPODArray) as soon as it holds one value, so one column per group costs ~4 KiB per +/// group regardless of k. That dominates the memory of a GROUP BY with millions of groups, while +/// these stores grow with the number of values actually retained. +template <typename T> +struct FixedTopKValueStore { + DorisVector<T> values; + + size_t size() const { return values.size(); } + void reserve(size_t n) { values.reserve(n); } + void clear() { values.clear(); } + const T& at(size_t slot) const { return values[slot]; } + void append(const T& value) { values.push_back(value); } + void swap(FixedTopKValueStore& other) noexcept { values.swap(other.values); } +}; + +/// A StringRef cannot be stored: the arena reallocates and would invalidate the pointers. A slot +/// keeps an offset/length into the bytes instead and materializes a view on access. +struct StringTopKValueStore { + struct Span { + uint32_t offset = 0; + uint32_t length = 0; + }; + + /// The initial allocation is small on purpose; see the comment on FixedTopKValueStore. + PODArray<char, 32> bytes; + DorisVector<Span> spans; + + size_t size() const { return spans.size(); } + void reserve(size_t n) { spans.reserve(n); } + void clear() { + bytes.clear(); + spans.clear(); + } + StringRef at(size_t slot) const { + const Span& span = spans[slot]; + return StringRef(bytes.data() + span.offset, span.length); + } + void append(StringRef value) { + DCHECK_LE(bytes.size() + value.size, std::numeric_limits<uint32_t>::max()); + DCHECK_LE(value.size, std::numeric_limits<uint32_t>::max()); + spans.push_back({static_cast<uint32_t>(bytes.size()), static_cast<uint32_t>(value.size)}); + if (value.size != 0) { + bytes.insert(value.data, value.data + value.size); + } + } + void swap(StringTopKValueStore& other) noexcept { + bytes.swap(other.bytes); + spans.swap(other.spans); + } +}; + +/** + * State shared by min_topk and max_topk, for one value type. + * + * The distinct values sit in slots (see the stores above), with one Slot record per value carrying + * the row count and the links of a treap ordered best-first (ascending values for MIN, descending + * for MAX). A row costs one O(log k) descent: it either finds its value and bumps the count or + * links a new slot at the leaf. A sorted vector would shift O(k) elements for the same insert. + * + * The treap stores explicit child links because, unlike a heap, its shape is not the complete tree + * that an array index can describe: rotations move entries, and the in-order sequence is the value + * order rather than the arrival order. A heap could evict and append in O(log k), but it cannot + * find an existing value without a hash index or a linear scan, and it cannot produce the sorted + * result without sorting. The ordered tree is the structure that keeps find, insert and eviction + * all at O(log k) with comparisons only. + * + * An entry that is evicted keeps its slot and value until the next compaction. The tree never + * points at a dead slot. + * + * With k == 0 the state holds at most the single extreme value and its row count; the caller + * derives truncation from the list length, so the state carries no separate bound. + */ +template <PrimitiveType T> +struct AggregateTopKCountData { + using CppType = typename PrimitiveTypeTraits<T>::CppType; + static constexpr bool IS_STRING = is_string_type(T); + using ValueStore = + std::conditional_t<IS_STRING, StringTopKValueStore, FixedTopKValueStore<CppType>>; + using ValueRef = std::conditional_t<IS_STRING, StringRef, const CppType&>; + + static constexpr int32_t NO_SLOT = -1; + // Row counts are never negative, so this cannot collide with a live entry. + static constexpr int64_t DEAD_COUNT = -1; + + /// One slot's metadata; the value itself lives in `values` at the same index. + struct Slot { + int64_t count = 0; + int32_t left = NO_SLOT; + int32_t right = NO_SLOT; + uint32_t priority = 0; + }; + + bool has_k = false; + int64_t k = 0; + // Indexed by slot; an evicted slot stays as a dead entry until compaction. + ValueStore values; + DorisVector<Slot> slots; + int32_t root = NO_SLOT; + // The last live slot in the aggregate's order: the value an incoming row has to beat to enter + // a full list. Cached so that rejecting a worse value costs one comparison, not a descent. + int32_t worst = NO_SLOT; + int64_t live = 0; + // Slots evicted since the last compaction. + int64_t dead = 0; + uint64_t rng = 0x243F6A8885A308D3ULL; + + AggregateTopKCountData() = default; + AggregateTopKCountData(const AggregateTopKCountData&) = delete; + AggregateTopKCountData& operator=(const AggregateTopKCountData&) = delete; +}; + +/** + * min_topk(x, k) / max_topk(x, k). + * + * Result: ARRAY<STRUCT<value T, count BIGINT>>. NULL inputs are ignored: add validates k for + * every row and then skips a NULL x, so the aggregate deliberately does not use the + * nullable-argument wrapper. The array is ordered by value (ascending for MIN), holds at most k + * distinct values with count > 0, and carries no truncation flag: the caller infers it from the + * length (a list of exactly k entries is treated as possibly truncated). + * + * With k = 0 the array holds at most the extreme value and its row count, which lets the + * append-only IVM path compute MIN/MAX from the delta extreme alone. + */ +template <PrimitiveType T, bool IsMin> +class AggregateFunctionTopK final + : public IAggregateFunctionDataHelper<AggregateTopKCountData<T>, + AggregateFunctionTopK<T, IsMin>, false>, + public MultiExpression, + public NotNullableAggregateFunction { +public: + using State = AggregateTopKCountData<T>; + using Base = IAggregateFunctionDataHelper<State, AggregateFunctionTopK<T, IsMin>, false>; + using Slot = typename State::Slot; + using ValueStore = typename State::ValueStore; + using ValueRef = typename State::ValueRef; + using ColumnType = typename PrimitiveTypeTraits<T>::ColumnType; + static constexpr bool IS_STRING = State::IS_STRING; + static constexpr int32_t NO_SLOT = State::NO_SLOT; + + explicit AggregateFunctionTopK(const DataTypes& argument_types_) : Base(argument_types_) { + DCHECK_EQ(argument_types_.size(), 2); + } + + String get_name() const override { return IsMin ? "min_topk" : "max_topk"; } + + DataTypePtr get_return_type() const override { + return make_return_type(remove_nullable(this->argument_types[0])); + } + + void add(AggregateDataPtr __restrict place, const IColumn** columns, ssize_t row_num, + Arena& arena) const override { + auto& state = Base::data(place); + // k is read for every row, including rows where x is NULL: an evaluated count that is + // NULL or out of range must fail the query instead of being skipped like a NULL value. + check_k(state, *columns[1], row_num); + const IColumn* value_column = columns[0]; + if (const auto* nullable = check_and_get_column<ColumnNullable>(value_column)) { + if (nullable->is_null_at(row_num)) { + return; + } + value_column = &nullable->get_nested_column(); + } + add_value(state, *value_column, row_num); + } + + void merge(AggregateDataPtr __restrict place, ConstAggregateDataPtr rhs, + Arena&) const override { + auto& state = Base::data(place); + const auto& other = Base::data(rhs); + // A partial state without has_k never saw any row, so it carries no information. + if (!other.has_k) { + return; + } + if (!state.has_k) { + state.has_k = true; + state.k = other.k; + } else if (UNLIKELY(state.k != other.k)) { + // Ordinary partial/final aggregation always shares the same constant k, but the + // public _state/_merge/_union combinators do not encode k in the AggState type, so a + // UNION ALL of states built with different literal k can reach this merge. + throw Exception(ErrorCode::INVALID_ARGUMENT, + "{} states have incompatible counts: {} vs {}", get_name(), state.k, + other.k); + } + if (state.k == 0) { + merge_extreme(state, other); + return; + } + merge_values(state, other); + } + + void serialize(ConstAggregateDataPtr __restrict place, BufferWritable& buf) const override { + const auto& state = Base::data(place); + buf.write_binary(state.has_k); + buf.write_binary(state.k); + // The slots are in arrival order and may include dead ones, while the serialized form is + // the sorted (value, count) list the reader rebuilds the tree from. + DorisVector<int32_t> order; + collect_in_order(state, order); + auto sorted_values = remove_nullable(this->argument_types[0])->create_column(); + sorted_values->reserve(order.size()); + DorisVector<int64_t> sorted_counts; + sorted_counts.reserve(order.size()); + for (int32_t slot : order) { + insert_value(*sorted_values, state.values.at(slot)); + sorted_counts.push_back(state.slots[slot].count); + } + auto size = static_cast<int64_t>(sorted_counts.size()); + buf.write_binary(size); + write_column(buf, *sorted_values); + if (size != 0) { + // An empty list keeps an empty vector, whose data() is null; passing it to memcpy is + // undefined even for a zero length. + buf.write(reinterpret_cast<const char*>(sorted_counts.data()), + sorted_counts.size() * sizeof(int64_t)); + } + } + + void deserialize(AggregateDataPtr __restrict place, BufferReadable& buf, + Arena&) const override { + auto& state = Base::data(place); + buf.read_binary(state.has_k); + buf.read_binary(state.k); + int64_t size = 0; + buf.read_binary(size); + auto values = read_column(buf); + DCHECK_EQ(static_cast<int64_t>(values->size()), size); + DorisVector<int64_t> counts(static_cast<size_t>(size)); + if (size != 0) { + buf.read(reinterpret_cast<char*>(counts.data()), size * sizeof(int64_t)); + } + ValueStore value_store; + value_store.reserve(static_cast<size_t>(size)); + for (int64_t row = 0; row != size; ++row) { + value_store.append(value_at(*values, static_cast<size_t>(row))); + } + assign_sorted(state, std::move(value_store), std::move(counts)); + } + + void insert_result_into(ConstAggregateDataPtr __restrict place, IColumn& to) const override { + const auto& state = Base::data(place); + auto& result_array = assert_cast<ColumnArray&>(to); + // The array element and both fields are nullable in the planned type, but a produced entry + // never is. + auto& pair_nullable = assert_cast<ColumnNullable&>(result_array.get_data()); + auto& pair_column = assert_cast<ColumnStruct&>(pair_nullable.get_nested_column()); + auto& value_column = assert_cast<ColumnNullable&>(pair_column.get_column(0)); + auto& count_column = assert_cast<ColumnNullable&>(pair_column.get_column(1)); + DorisVector<int32_t> order; + collect_in_order(state, order); + for (int32_t slot : order) { + insert_value(value_column.get_nested_column(), state.values.at(slot)); + value_column.get_null_map_data().push_back(0); + assert_cast<ColumnInt64&>(count_column.get_nested_column()) + .get_data() + .push_back(state.slots[slot].count); + count_column.get_null_map_data().push_back(0); + pair_nullable.get_null_map_data().push_back(0); + } + result_array.get_offsets().push_back(pair_column.size()); + } + + static DataTypePtr make_return_type(const DataTypePtr& value_type) { + // Both fields are nullable in the planned type: a consumer reads a field through + // element_at, which is always nullable, and named_struct then records nullable fields. + // Produced entries never contain a NULL field. + auto pair_type = std::make_shared<DataTypeStruct>( + DataTypes {make_nullable(value_type), + make_nullable(std::make_shared<DataTypeInt64>())}, + Strings {"value", "count"}); + // ArrayType defaults contains_null to true on the FE side, so the pair element is nullable + // in the planned type even though the aggregate never produces a NULL entry. + return std::make_shared<DataTypeArray>(make_nullable(pair_type)); + } + +private: + /// The value of the row in `column`. The returned view is valid while the column is alive. + static ValueRef value_at(const IColumn& column, size_t row_num) { + if constexpr (IS_STRING) { + return assert_cast<const ColumnType&, TypeCheckOnRelease::DISABLE>(column).get_data_at( + row_num); + } else { + return assert_cast<const ColumnType&, TypeCheckOnRelease::DISABLE>(column) + .get_data()[row_num]; + } + } + + /// Exactly the comparison the value column would apply: CompareFloat for FLOAT/DOUBLE through + /// Compare::compare (the column uses the same helper), the same memcmp helper ColumnString + /// uses for strings, and the column's scale-independent ordering for every other type. + static int compare_values(ValueRef lhs, ValueRef rhs) { + if constexpr (IS_STRING) { + return memcmp_small_allow_overflow15( + reinterpret_cast<const uint8_t*>(lhs.data), lhs.size, + reinterpret_cast<const uint8_t*>(rhs.data), rhs.size); + } else { + return Compare::compare(lhs, rhs); + } + } + + static void insert_value(IColumn& column, ValueRef value) { + if constexpr (IS_STRING) { + assert_cast<ColumnString&>(column).insert_data(value.data, value.size); + } else { + assert_cast<ColumnType&>(column).get_data().push_back(value); + } + } + + /// True when `cmp` puts the left operand before the right one in the aggregate's order. + static bool comes_before(int cmp) { return IsMin ? cmp < 0 : cmp > 0; } + + /// True when `cmp`, which compares a value with the worst retained one, puts it after it: the + /// value is worse and cannot enter the list. + static bool is_worse(int cmp) { return IsMin ? cmp > 0 : cmp < 0; } + + /// splitmix64. The priorities only have to be independent enough to keep the treap balanced; + /// a per-state sequence keeps the shape reproducible for the same insertion order. + static uint32_t next_priority(State& state) { + uint64_t z = (state.rng += 0x9E3779B97F4A7C15ULL); + z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL; + z = (z ^ (z >> 27)) * 0x94D049BB133111EBULL; + return static_cast<uint32_t>(z ^ (z >> 31)); + } + + static int compare_slots(const State& state, int32_t lhs, int32_t rhs) { + return compare_values(state.values.at(lhs), state.values.at(rhs)); + } + + /// Appends (value, count) as the next slot. + static int32_t append_slot(State& state, ValueRef value, int64_t count) { + auto slot = static_cast<int32_t>(state.slots.size()); + DCHECK_LT(slot, std::numeric_limits<int32_t>::max()); + state.values.append(value); + state.slots.push_back(Slot {count, State::NO_SLOT, State::NO_SLOT, next_priority(state)}); + ++state.live; + return slot; + } + + static int32_t append_slot(State& state, const IColumn& column, size_t row_num, int64_t count) { + return append_slot(state, value_at(column, row_num), count); + } + + /// Rotates the subtree rooted at `node` right: the left child becomes the root and `node` its + /// right child. The in-order sequence is unchanged, so the treap stays ordered. + static int32_t rotate_right(State& state, int32_t node) { + int32_t pivot = state.slots[node].left; + state.slots[node].left = state.slots[pivot].right; + state.slots[pivot].right = node; + return pivot; + } + + static int32_t rotate_left(State& state, int32_t node) { + int32_t pivot = state.slots[node].right; + state.slots[node].right = state.slots[pivot].left; + state.slots[pivot].left = node; + return pivot; + } + + /// One descent that performs the whole per-row update: an equal value bumps its count and + /// reports no insertion; a missing value is linked as a new slot at the leaf and rotated up + /// while the recursion unwinds. `inserted_slot` receives the new slot, NO_SLOT on a hit. + static int32_t add_or_increment(State& state, int32_t node, const IColumn& column, + size_t row_num, int32_t* inserted_slot) { + if (node == State::NO_SLOT) { + *inserted_slot = append_slot(state, column, row_num, 1); + return *inserted_slot; + } + int cmp = compare_values(state.values.at(node), value_at(column, row_num)); + if (cmp == 0) { + state.slots[node].count = checked_add(state.slots[node].count, 1); + return node; + } + if (comes_before(cmp)) { + // The stored value is the better one, so the input belongs in the right subtree. + state.slots[node].right = add_or_increment(state, state.slots[node].right, column, + row_num, inserted_slot); + if (state.slots[state.slots[node].right].priority > state.slots[node].priority) { + return rotate_left(state, node); + } + } else { + state.slots[node].left = + add_or_increment(state, state.slots[node].left, column, row_num, inserted_slot); + if (state.slots[state.slots[node].left].priority > state.slots[node].priority) { + return rotate_right(state, node); + } + } + return node; + } + + /// Removes the last entry in the aggregate's order (the largest value for MIN) and returns the + /// new subtree root. That entry has no right child, so its left subtree takes its place; the + /// child's priority is below the parent's, so no rotation is needed. + static int32_t erase_worst(State& state, int32_t node, int32_t* removed) { + if (state.slots[node].right == State::NO_SLOT) { + *removed = node; + return state.slots[node].left; + } + state.slots[node].right = erase_worst(state, state.slots[node].right, removed); + return node; + } + + static void refresh_worst(State& state) { + int32_t node = state.root; + while (node != State::NO_SLOT && state.slots[node].right != State::NO_SLOT) { + node = state.slots[node].right; + } + state.worst = node; + } + + /// Slots in the aggregate's order, i.e. the order of the result array. + static void collect_in_order(const State& state, DorisVector<int32_t>& slots) { + slots.reserve(static_cast<size_t>(state.live)); + DorisVector<int32_t> stack; + stack.reserve(64); + int32_t node = state.root; + while (node != State::NO_SLOT || !stack.empty()) { + while (node != State::NO_SLOT) { + stack.push_back(node); + node = state.slots[node].left; + } + node = stack.back(); + stack.pop_back(); + // Only live slots are reachable from the root; an evicted slot has no row count. + DCHECK_GE(state.slots[node].count, 0); + slots.push_back(node); + node = state.slots[node].right; + } + DCHECK_EQ(static_cast<int64_t>(slots.size()), state.live); + } + + /// Rebuilds the tree from `values`, already in the aggregate's order, and the matching counts. + /// The treap is built in O(size) as a Cartesian tree over the fresh priorities. + static void assign_sorted(State& state, ValueStore values, DorisVector<int64_t> counts) { + const size_t size = counts.size(); + DCHECK_EQ(values.size(), size); + state.values.swap(values); + state.slots.assign(size, Slot {}); + for (size_t slot = 0; slot != size; ++slot) { + state.slots[slot].count = counts[slot]; + } + state.live = static_cast<int64_t>(size); + state.dead = 0; + DorisVector<int32_t> stack; + stack.reserve(size); + for (size_t slot = 0; slot != size; ++slot) { + state.slots[slot].priority = next_priority(state); + int32_t last = State::NO_SLOT; + while (!stack.empty() && + state.slots[stack.back()].priority < state.slots[slot].priority) { + last = stack.back(); + stack.pop_back(); + } + state.slots[slot].left = last; + if (!stack.empty()) { + state.slots[stack.back()].right = static_cast<int32_t>(slot); + } + stack.push_back(static_cast<int32_t>(slot)); + } + state.root = stack.empty() ? State::NO_SLOT : stack.front(); + refresh_worst(state); + } + + /// Reclaims the slots evicted since the last compaction. The O(k) rebuild is charged to the + /// k/2 evictions that produced the dead slots, so it stays O(1) amortized per row. + static void compact(State& state) { + DorisVector<int32_t> slots; + collect_in_order(state, slots); + ValueStore values; + values.reserve(slots.size()); + DorisVector<int64_t> counts; + counts.reserve(slots.size()); + for (int32_t slot : slots) { + values.append(state.values.at(slot)); + counts.push_back(state.slots[slot].count); + } + assign_sorted(state, std::move(values), std::move(counts)); + } + + static void check_k(State& state, const IColumn& count_column, size_t row_num) { + const IColumn* nested = &count_column; + if (const auto* nullable = check_and_get_column<ColumnNullable>(&count_column)) { + if (UNLIKELY(nullable->is_null_at(row_num))) { + throw Exception(ErrorCode::INVALID_ARGUMENT, + "the count of min_topk and max_topk must be a non-null " + "constant"); + } + nested = &nullable->get_nested_column(); + } + int64_t value = + assert_cast<const ColumnInt64&, TypeCheckOnRelease::DISABLE>(*nested).get_element( + row_num); + if (!state.has_k) { + if (UNLIKELY(value < 0)) { + throw Exception(ErrorCode::INVALID_ARGUMENT, + "the count of min_topk and max_topk must be a non-negative " + "constant, but got {}", + value); + } + state.has_k = true; + state.k = value; + } else { + // k comes from the same constant expression on every row and every partial state. + DCHECK_EQ(state.k, value); + } + } + + static int64_t checked_add(int64_t lhs, int64_t rhs) { + int64_t result = 0; + if (UNLIKELY(__builtin_add_overflow(lhs, rhs, &result))) { + throw Exception(ErrorCode::INVALID_ARGUMENT, "min_topk/max_topk value count overflow"); + } + return result; + } + + /// k == 0 fast path: only the extreme value is tracked. + static void add_value(State& state, const IColumn& column, size_t row_num) { + if (state.k == 0) { + if (state.live == 0) { + state.root = append_slot(state, column, row_num, 1); + state.worst = state.root; + return; + } + int cmp = compare_values(value_at(column, row_num), state.values.at(0)); + if (cmp == 0) { + state.slots[0].count = checked_add(state.slots[0].count, 1); + } else if (comes_before(cmp)) { + // The new value is better: replace the tracked extreme. The slot and the + // single-node tree stay in place, only the stored value changes. + state.values.clear(); + state.values.append(value_at(column, row_num)); + state.slots[0].count = 1; + } + return; + } + + // Dead slots are reclaimed before they can outnumber the live ones; see compact(). + if (state.dead > state.live) { Review Comment: [P2] Reclaim evicted strings according to their byte size Fill a MIN state at `k=1024` with 1024 distinct 1 MiB `z`-prefixed strings, then replace them with 1024 distinct short `a`-prefixed strings. The live result is a few KiB, but `dead==live==1024`, so this condition never compacts, including on subsequent duplicate or rejected rows. At least 1 GiB of discarded payload remains owned until group destruction or a later eviction; processing many groups sequentially can exhaust query memory despite tiny final states. Track discarded/live bytes and trigger byte-aware string reclamation. This is separate from offset wrap: it occurs below 4 GiB and remains after changing offsets to a wider type. ########## be/src/exprs/aggregate/aggregate_function_topk.h: ########## @@ -0,0 +1,741 @@ +// 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. + +#pragma once + +#include <algorithm> +#include <cstdint> +#include <limits> +#include <memory> +#include <type_traits> +#include <utility> + +#include "common/check.h" +#include "common/compare.h" +#include "common/exception.h" +#include "core/assert_cast.h" +#include "core/column/column.h" +#include "core/column/column_array.h" +#include "core/column/column_decimal.h" +#include "core/column/column_nullable.h" +#include "core/column/column_string.h" +#include "core/column/column_struct.h" +#include "core/column/column_vector.h" +#include "core/custom_allocator.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 "core/data_type/data_type_struct.h" +#include "core/data_type/primitive_type.h" +#include "core/memcmp_small.h" +#include "core/pod_array.h" +#include "core/string_buffer.hpp" +#include "core/string_ref.h" +#include "core/types.h" +#include "exprs/aggregate/aggregate_function.h" + +namespace doris { + +/// The values of one group, in arrival order, one entry per slot. +/// +/// They deliberately do not live in a Doris column: a column's container allocates a 4 KiB initial +/// chunk (PaddedPODArray) as soon as it holds one value, so one column per group costs ~4 KiB per +/// group regardless of k. That dominates the memory of a GROUP BY with millions of groups, while +/// these stores grow with the number of values actually retained. +template <typename T> +struct FixedTopKValueStore { + DorisVector<T> values; + + size_t size() const { return values.size(); } + void reserve(size_t n) { values.reserve(n); } + void clear() { values.clear(); } + const T& at(size_t slot) const { return values[slot]; } + void append(const T& value) { values.push_back(value); } + void swap(FixedTopKValueStore& other) noexcept { values.swap(other.values); } +}; + +/// A StringRef cannot be stored: the arena reallocates and would invalidate the pointers. A slot +/// keeps an offset/length into the bytes instead and materializes a view on access. +struct StringTopKValueStore { + struct Span { + uint32_t offset = 0; + uint32_t length = 0; + }; + + /// The initial allocation is small on purpose; see the comment on FixedTopKValueStore. + PODArray<char, 32> bytes; + DorisVector<Span> spans; + + size_t size() const { return spans.size(); } + void reserve(size_t n) { spans.reserve(n); } + void clear() { + bytes.clear(); + spans.clear(); + } + StringRef at(size_t slot) const { + const Span& span = spans[slot]; + return StringRef(bytes.data() + span.offset, span.length); + } + void append(StringRef value) { + DCHECK_LE(bytes.size() + value.size, std::numeric_limits<uint32_t>::max()); + DCHECK_LE(value.size, std::numeric_limits<uint32_t>::max()); + spans.push_back({static_cast<uint32_t>(bytes.size()), static_cast<uint32_t>(value.size)}); + if (value.size != 0) { + bytes.insert(value.data, value.data + value.size); + } + } + void swap(StringTopKValueStore& other) noexcept { + bytes.swap(other.bytes); + spans.swap(other.spans); + } +}; + +/** + * State shared by min_topk and max_topk, for one value type. + * + * The distinct values sit in slots (see the stores above), with one Slot record per value carrying + * the row count and the links of a treap ordered best-first (ascending values for MIN, descending + * for MAX). A row costs one O(log k) descent: it either finds its value and bumps the count or + * links a new slot at the leaf. A sorted vector would shift O(k) elements for the same insert. + * + * The treap stores explicit child links because, unlike a heap, its shape is not the complete tree + * that an array index can describe: rotations move entries, and the in-order sequence is the value + * order rather than the arrival order. A heap could evict and append in O(log k), but it cannot + * find an existing value without a hash index or a linear scan, and it cannot produce the sorted + * result without sorting. The ordered tree is the structure that keeps find, insert and eviction + * all at O(log k) with comparisons only. + * + * An entry that is evicted keeps its slot and value until the next compaction. The tree never + * points at a dead slot. + * + * With k == 0 the state holds at most the single extreme value and its row count; the caller + * derives truncation from the list length, so the state carries no separate bound. + */ +template <PrimitiveType T> +struct AggregateTopKCountData { + using CppType = typename PrimitiveTypeTraits<T>::CppType; + static constexpr bool IS_STRING = is_string_type(T); + using ValueStore = + std::conditional_t<IS_STRING, StringTopKValueStore, FixedTopKValueStore<CppType>>; + using ValueRef = std::conditional_t<IS_STRING, StringRef, const CppType&>; + + static constexpr int32_t NO_SLOT = -1; + // Row counts are never negative, so this cannot collide with a live entry. + static constexpr int64_t DEAD_COUNT = -1; + + /// One slot's metadata; the value itself lives in `values` at the same index. + struct Slot { + int64_t count = 0; + int32_t left = NO_SLOT; + int32_t right = NO_SLOT; + uint32_t priority = 0; + }; + + bool has_k = false; + int64_t k = 0; + // Indexed by slot; an evicted slot stays as a dead entry until compaction. + ValueStore values; + DorisVector<Slot> slots; + int32_t root = NO_SLOT; + // The last live slot in the aggregate's order: the value an incoming row has to beat to enter + // a full list. Cached so that rejecting a worse value costs one comparison, not a descent. + int32_t worst = NO_SLOT; + int64_t live = 0; + // Slots evicted since the last compaction. + int64_t dead = 0; + uint64_t rng = 0x243F6A8885A308D3ULL; + + AggregateTopKCountData() = default; + AggregateTopKCountData(const AggregateTopKCountData&) = delete; + AggregateTopKCountData& operator=(const AggregateTopKCountData&) = delete; +}; + +/** + * min_topk(x, k) / max_topk(x, k). + * + * Result: ARRAY<STRUCT<value T, count BIGINT>>. NULL inputs are ignored: add validates k for + * every row and then skips a NULL x, so the aggregate deliberately does not use the + * nullable-argument wrapper. The array is ordered by value (ascending for MIN), holds at most k + * distinct values with count > 0, and carries no truncation flag: the caller infers it from the + * length (a list of exactly k entries is treated as possibly truncated). + * + * With k = 0 the array holds at most the extreme value and its row count, which lets the + * append-only IVM path compute MIN/MAX from the delta extreme alone. + */ +template <PrimitiveType T, bool IsMin> +class AggregateFunctionTopK final + : public IAggregateFunctionDataHelper<AggregateTopKCountData<T>, + AggregateFunctionTopK<T, IsMin>, false>, + public MultiExpression, + public NotNullableAggregateFunction { +public: + using State = AggregateTopKCountData<T>; + using Base = IAggregateFunctionDataHelper<State, AggregateFunctionTopK<T, IsMin>, false>; + using Slot = typename State::Slot; + using ValueStore = typename State::ValueStore; + using ValueRef = typename State::ValueRef; + using ColumnType = typename PrimitiveTypeTraits<T>::ColumnType; + static constexpr bool IS_STRING = State::IS_STRING; + static constexpr int32_t NO_SLOT = State::NO_SLOT; + + explicit AggregateFunctionTopK(const DataTypes& argument_types_) : Base(argument_types_) { + DCHECK_EQ(argument_types_.size(), 2); + } + + String get_name() const override { return IsMin ? "min_topk" : "max_topk"; } + + DataTypePtr get_return_type() const override { + return make_return_type(remove_nullable(this->argument_types[0])); + } + + void add(AggregateDataPtr __restrict place, const IColumn** columns, ssize_t row_num, + Arena& arena) const override { + auto& state = Base::data(place); + // k is read for every row, including rows where x is NULL: an evaluated count that is + // NULL or out of range must fail the query instead of being skipped like a NULL value. + check_k(state, *columns[1], row_num); + const IColumn* value_column = columns[0]; + if (const auto* nullable = check_and_get_column<ColumnNullable>(value_column)) { + if (nullable->is_null_at(row_num)) { + return; + } + value_column = &nullable->get_nested_column(); + } + add_value(state, *value_column, row_num); + } + + void merge(AggregateDataPtr __restrict place, ConstAggregateDataPtr rhs, + Arena&) const override { + auto& state = Base::data(place); + const auto& other = Base::data(rhs); + // A partial state without has_k never saw any row, so it carries no information. + if (!other.has_k) { + return; + } + if (!state.has_k) { + state.has_k = true; + state.k = other.k; + } else if (UNLIKELY(state.k != other.k)) { + // Ordinary partial/final aggregation always shares the same constant k, but the + // public _state/_merge/_union combinators do not encode k in the AggState type, so a + // UNION ALL of states built with different literal k can reach this merge. + throw Exception(ErrorCode::INVALID_ARGUMENT, + "{} states have incompatible counts: {} vs {}", get_name(), state.k, + other.k); + } + if (state.k == 0) { + merge_extreme(state, other); + return; + } + merge_values(state, other); + } + + void serialize(ConstAggregateDataPtr __restrict place, BufferWritable& buf) const override { + const auto& state = Base::data(place); + buf.write_binary(state.has_k); + buf.write_binary(state.k); + // The slots are in arrival order and may include dead ones, while the serialized form is + // the sorted (value, count) list the reader rebuilds the tree from. + DorisVector<int32_t> order; + collect_in_order(state, order); + auto sorted_values = remove_nullable(this->argument_types[0])->create_column(); + sorted_values->reserve(order.size()); + DorisVector<int64_t> sorted_counts; + sorted_counts.reserve(order.size()); + for (int32_t slot : order) { + insert_value(*sorted_values, state.values.at(slot)); + sorted_counts.push_back(state.slots[slot].count); + } + auto size = static_cast<int64_t>(sorted_counts.size()); + buf.write_binary(size); + write_column(buf, *sorted_values); + if (size != 0) { + // An empty list keeps an empty vector, whose data() is null; passing it to memcpy is + // undefined even for a zero length. + buf.write(reinterpret_cast<const char*>(sorted_counts.data()), + sorted_counts.size() * sizeof(int64_t)); + } + } + + void deserialize(AggregateDataPtr __restrict place, BufferReadable& buf, + Arena&) const override { + auto& state = Base::data(place); + buf.read_binary(state.has_k); + buf.read_binary(state.k); + int64_t size = 0; + buf.read_binary(size); + auto values = read_column(buf); + DCHECK_EQ(static_cast<int64_t>(values->size()), size); + DorisVector<int64_t> counts(static_cast<size_t>(size)); + if (size != 0) { + buf.read(reinterpret_cast<char*>(counts.data()), size * sizeof(int64_t)); + } + ValueStore value_store; + value_store.reserve(static_cast<size_t>(size)); + for (int64_t row = 0; row != size; ++row) { + value_store.append(value_at(*values, static_cast<size_t>(row))); + } + assign_sorted(state, std::move(value_store), std::move(counts)); + } + + void insert_result_into(ConstAggregateDataPtr __restrict place, IColumn& to) const override { + const auto& state = Base::data(place); + auto& result_array = assert_cast<ColumnArray&>(to); + // The array element and both fields are nullable in the planned type, but a produced entry + // never is. + auto& pair_nullable = assert_cast<ColumnNullable&>(result_array.get_data()); + auto& pair_column = assert_cast<ColumnStruct&>(pair_nullable.get_nested_column()); + auto& value_column = assert_cast<ColumnNullable&>(pair_column.get_column(0)); + auto& count_column = assert_cast<ColumnNullable&>(pair_column.get_column(1)); + DorisVector<int32_t> order; + collect_in_order(state, order); + for (int32_t slot : order) { + insert_value(value_column.get_nested_column(), state.values.at(slot)); + value_column.get_null_map_data().push_back(0); + assert_cast<ColumnInt64&>(count_column.get_nested_column()) + .get_data() + .push_back(state.slots[slot].count); + count_column.get_null_map_data().push_back(0); + pair_nullable.get_null_map_data().push_back(0); + } + result_array.get_offsets().push_back(pair_column.size()); + } + + static DataTypePtr make_return_type(const DataTypePtr& value_type) { + // Both fields are nullable in the planned type: a consumer reads a field through + // element_at, which is always nullable, and named_struct then records nullable fields. + // Produced entries never contain a NULL field. + auto pair_type = std::make_shared<DataTypeStruct>( + DataTypes {make_nullable(value_type), + make_nullable(std::make_shared<DataTypeInt64>())}, + Strings {"value", "count"}); + // ArrayType defaults contains_null to true on the FE side, so the pair element is nullable + // in the planned type even though the aggregate never produces a NULL entry. + return std::make_shared<DataTypeArray>(make_nullable(pair_type)); + } + +private: + /// The value of the row in `column`. The returned view is valid while the column is alive. + static ValueRef value_at(const IColumn& column, size_t row_num) { + if constexpr (IS_STRING) { + return assert_cast<const ColumnType&, TypeCheckOnRelease::DISABLE>(column).get_data_at( + row_num); + } else { + return assert_cast<const ColumnType&, TypeCheckOnRelease::DISABLE>(column) + .get_data()[row_num]; + } + } + + /// Exactly the comparison the value column would apply: CompareFloat for FLOAT/DOUBLE through + /// Compare::compare (the column uses the same helper), the same memcmp helper ColumnString + /// uses for strings, and the column's scale-independent ordering for every other type. + static int compare_values(ValueRef lhs, ValueRef rhs) { + if constexpr (IS_STRING) { + return memcmp_small_allow_overflow15( + reinterpret_cast<const uint8_t*>(lhs.data), lhs.size, + reinterpret_cast<const uint8_t*>(rhs.data), rhs.size); + } else { + return Compare::compare(lhs, rhs); + } + } + + static void insert_value(IColumn& column, ValueRef value) { + if constexpr (IS_STRING) { + assert_cast<ColumnString&>(column).insert_data(value.data, value.size); + } else { + assert_cast<ColumnType&>(column).get_data().push_back(value); + } + } + + /// True when `cmp` puts the left operand before the right one in the aggregate's order. + static bool comes_before(int cmp) { return IsMin ? cmp < 0 : cmp > 0; } + + /// True when `cmp`, which compares a value with the worst retained one, puts it after it: the + /// value is worse and cannot enter the list. + static bool is_worse(int cmp) { return IsMin ? cmp > 0 : cmp < 0; } + + /// splitmix64. The priorities only have to be independent enough to keep the treap balanced; + /// a per-state sequence keeps the shape reproducible for the same insertion order. + static uint32_t next_priority(State& state) { + uint64_t z = (state.rng += 0x9E3779B97F4A7C15ULL); + z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL; + z = (z ^ (z >> 27)) * 0x94D049BB133111EBULL; + return static_cast<uint32_t>(z ^ (z >> 31)); + } + + static int compare_slots(const State& state, int32_t lhs, int32_t rhs) { + return compare_values(state.values.at(lhs), state.values.at(rhs)); + } + + /// Appends (value, count) as the next slot. + static int32_t append_slot(State& state, ValueRef value, int64_t count) { + auto slot = static_cast<int32_t>(state.slots.size()); + DCHECK_LT(slot, std::numeric_limits<int32_t>::max()); + state.values.append(value); + state.slots.push_back(Slot {count, State::NO_SLOT, State::NO_SLOT, next_priority(state)}); + ++state.live; + return slot; + } + + static int32_t append_slot(State& state, const IColumn& column, size_t row_num, int64_t count) { + return append_slot(state, value_at(column, row_num), count); + } + + /// Rotates the subtree rooted at `node` right: the left child becomes the root and `node` its + /// right child. The in-order sequence is unchanged, so the treap stays ordered. + static int32_t rotate_right(State& state, int32_t node) { + int32_t pivot = state.slots[node].left; + state.slots[node].left = state.slots[pivot].right; + state.slots[pivot].right = node; + return pivot; + } + + static int32_t rotate_left(State& state, int32_t node) { + int32_t pivot = state.slots[node].right; + state.slots[node].right = state.slots[pivot].left; + state.slots[pivot].left = node; + return pivot; + } + + /// One descent that performs the whole per-row update: an equal value bumps its count and + /// reports no insertion; a missing value is linked as a new slot at the leaf and rotated up + /// while the recursion unwinds. `inserted_slot` receives the new slot, NO_SLOT on a hit. + static int32_t add_or_increment(State& state, int32_t node, const IColumn& column, + size_t row_num, int32_t* inserted_slot) { + if (node == State::NO_SLOT) { + *inserted_slot = append_slot(state, column, row_num, 1); + return *inserted_slot; + } + int cmp = compare_values(state.values.at(node), value_at(column, row_num)); + if (cmp == 0) { + state.slots[node].count = checked_add(state.slots[node].count, 1); + return node; + } + if (comes_before(cmp)) { + // The stored value is the better one, so the input belongs in the right subtree. + state.slots[node].right = add_or_increment(state, state.slots[node].right, column, + row_num, inserted_slot); + if (state.slots[state.slots[node].right].priority > state.slots[node].priority) { + return rotate_left(state, node); + } + } else { + state.slots[node].left = + add_or_increment(state, state.slots[node].left, column, row_num, inserted_slot); + if (state.slots[state.slots[node].left].priority > state.slots[node].priority) { + return rotate_right(state, node); + } + } + return node; + } + + /// Removes the last entry in the aggregate's order (the largest value for MIN) and returns the + /// new subtree root. That entry has no right child, so its left subtree takes its place; the + /// child's priority is below the parent's, so no rotation is needed. + static int32_t erase_worst(State& state, int32_t node, int32_t* removed) { + if (state.slots[node].right == State::NO_SLOT) { + *removed = node; + return state.slots[node].left; + } + state.slots[node].right = erase_worst(state, state.slots[node].right, removed); + return node; + } + + static void refresh_worst(State& state) { + int32_t node = state.root; + while (node != State::NO_SLOT && state.slots[node].right != State::NO_SLOT) { + node = state.slots[node].right; + } + state.worst = node; + } + + /// Slots in the aggregate's order, i.e. the order of the result array. + static void collect_in_order(const State& state, DorisVector<int32_t>& slots) { + slots.reserve(static_cast<size_t>(state.live)); + DorisVector<int32_t> stack; + stack.reserve(64); + int32_t node = state.root; + while (node != State::NO_SLOT || !stack.empty()) { + while (node != State::NO_SLOT) { + stack.push_back(node); + node = state.slots[node].left; + } + node = stack.back(); + stack.pop_back(); + // Only live slots are reachable from the root; an evicted slot has no row count. + DCHECK_GE(state.slots[node].count, 0); + slots.push_back(node); + node = state.slots[node].right; + } + DCHECK_EQ(static_cast<int64_t>(slots.size()), state.live); + } + + /// Rebuilds the tree from `values`, already in the aggregate's order, and the matching counts. + /// The treap is built in O(size) as a Cartesian tree over the fresh priorities. + static void assign_sorted(State& state, ValueStore values, DorisVector<int64_t> counts) { + const size_t size = counts.size(); + DCHECK_EQ(values.size(), size); + state.values.swap(values); + state.slots.assign(size, Slot {}); + for (size_t slot = 0; slot != size; ++slot) { + state.slots[slot].count = counts[slot]; + } + state.live = static_cast<int64_t>(size); + state.dead = 0; + DorisVector<int32_t> stack; + stack.reserve(size); + for (size_t slot = 0; slot != size; ++slot) { + state.slots[slot].priority = next_priority(state); + int32_t last = State::NO_SLOT; + while (!stack.empty() && + state.slots[stack.back()].priority < state.slots[slot].priority) { + last = stack.back(); + stack.pop_back(); + } + state.slots[slot].left = last; + if (!stack.empty()) { + state.slots[stack.back()].right = static_cast<int32_t>(slot); + } + stack.push_back(static_cast<int32_t>(slot)); + } + state.root = stack.empty() ? State::NO_SLOT : stack.front(); + refresh_worst(state); + } + + /// Reclaims the slots evicted since the last compaction. The O(k) rebuild is charged to the + /// k/2 evictions that produced the dead slots, so it stays O(1) amortized per row. + static void compact(State& state) { + DorisVector<int32_t> slots; + collect_in_order(state, slots); + ValueStore values; + values.reserve(slots.size()); + DorisVector<int64_t> counts; + counts.reserve(slots.size()); + for (int32_t slot : slots) { + values.append(state.values.at(slot)); + counts.push_back(state.slots[slot].count); + } + assign_sorted(state, std::move(values), std::move(counts)); + } + + static void check_k(State& state, const IColumn& count_column, size_t row_num) { + const IColumn* nested = &count_column; + if (const auto* nullable = check_and_get_column<ColumnNullable>(&count_column)) { + if (UNLIKELY(nullable->is_null_at(row_num))) { + throw Exception(ErrorCode::INVALID_ARGUMENT, + "the count of min_topk and max_topk must be a non-null " + "constant"); + } + nested = &nullable->get_nested_column(); + } + int64_t value = + assert_cast<const ColumnInt64&, TypeCheckOnRelease::DISABLE>(*nested).get_element( + row_num); + if (!state.has_k) { + if (UNLIKELY(value < 0)) { + throw Exception(ErrorCode::INVALID_ARGUMENT, + "the count of min_topk and max_topk must be a non-negative " + "constant, but got {}", + value); + } + state.has_k = true; + state.k = value; + } else { + // k comes from the same constant expression on every row and every partial state. + DCHECK_EQ(state.k, value); + } + } + + static int64_t checked_add(int64_t lhs, int64_t rhs) { + int64_t result = 0; + if (UNLIKELY(__builtin_add_overflow(lhs, rhs, &result))) { + throw Exception(ErrorCode::INVALID_ARGUMENT, "min_topk/max_topk value count overflow"); + } + return result; + } + + /// k == 0 fast path: only the extreme value is tracked. + static void add_value(State& state, const IColumn& column, size_t row_num) { + if (state.k == 0) { + if (state.live == 0) { + state.root = append_slot(state, column, row_num, 1); + state.worst = state.root; + return; + } + int cmp = compare_values(value_at(column, row_num), state.values.at(0)); + if (cmp == 0) { + state.slots[0].count = checked_add(state.slots[0].count, 1); + } else if (comes_before(cmp)) { + // The new value is better: replace the tracked extreme. The slot and the + // single-node tree stay in place, only the stored value changes. + state.values.clear(); + state.values.append(value_at(column, row_num)); + state.slots[0].count = 1; + } + return; + } + + // Dead slots are reclaimed before they can outnumber the live ones; see compact(). + if (state.dead > state.live) { + compact(state); + } + + if (state.live >= state.k) { + // A full list rejects or absorbs most rows: comparing with the eviction candidate + // costs one comparison, while the tree descent would cost log2(k). + int cmp = compare_values(value_at(column, row_num), state.values.at(state.worst)); + if (cmp == 0) { + state.slots[state.worst].count = checked_add(state.slots[state.worst].count, 1); + return; + } + if (is_worse(cmp)) { + return; + } + } + + int32_t inserted = State::NO_SLOT; + state.root = add_or_increment(state, state.root, column, row_num, &inserted); + if (inserted == State::NO_SLOT) { + return; + } + if (state.live <= state.k) { + // Not over capacity: the new entry competes with the cached worst for the eviction + // candidate role. + if (state.worst == State::NO_SLOT || + is_worse(compare_slots(state, inserted, state.worst))) { + state.worst = inserted; + } + return; + } + // The list was full and the value is new. It beat the cached worst above, so it cannot be + // the worst itself; the old worst is still the rightmost entry and leaves the list. + int32_t removed = State::NO_SLOT; + state.root = erase_worst(state, state.root, &removed); + DCHECK_EQ(removed, state.worst); + state.slots[removed].count = State::DEAD_COUNT; + --state.live; + ++state.dead; + refresh_worst(state); + } + + /// k == 0: the state holds at most the extreme value and its count. + static void merge_extreme(State& state, const State& other) { + if (other.live == 0) { + return; + } + if (state.live == 0) { + state.root = append_slot(state, other.values.at(0), other.slots[0].count); + state.worst = state.root; + return; + } + int cmp = compare_values(state.values.at(0), other.values.at(0)); + if (cmp == 0) { + state.slots[0].count = checked_add(state.slots[0].count, other.slots[0].count); + } else if (is_worse(cmp)) { + state.values.clear(); + state.values.append(other.values.at(0)); + state.slots[0].count = other.slots[0].count; + } + } + + /// Merges two ordered lists, summing the counts of equal values and keeping the k best. + /// + /// The merge stops as soon as k entries are produced. Both inputs are best-first and hold + /// distinct values, and equal values have already been combined, so every remaining value is + /// worse than all retained ones: it can neither enter the result nor overflow it. Stopping + /// here also avoids materializing a tail that would be truncated away. + static void merge_values(State& state, const State& other) { + if (other.live == 0) { + // The other partial state saw no non-NULL row. + return; + } + DorisVector<int32_t> lhs_slots; + DorisVector<int32_t> rhs_slots; + collect_in_order(state, lhs_slots); Review Comment: [P2] Skip partial states that cannot change a full result `min_topk_merge(min_topk_state(x,65536))` serializes one singleton per row, and `AggregateStateUnion` invokes this merge for each singleton. After MIN retains `0..65535`, every higher singleton still traverses/copies all 65536 retained values and `assign_sorted` rebuilds their tree. One million ascending inputs therefore cause over 61 billion unchanged-value copies after the list fills (source-derived, not benchmarked); MAX/descending input is symmetric. This serialized merge bypasses the existing `add_value` rejection fast path, and stopping after k outputs still rebuilds the unchanged prefix. Reject an entirely worse RHS using the cached worst and RHS best before collecting/rebuilding the prefix. -- 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]
