Skip to content

Commit

Permalink
[Enhancement] Support aggregate function map_agg.
Browse files Browse the repository at this point in the history
Signed-off-by: Song Jiacheng <[email protected]>
  • Loading branch information
Jcnessss committed Oct 17, 2024
1 parent 3c11722 commit 644a6f9
Show file tree
Hide file tree
Showing 6 changed files with 476 additions and 0 deletions.
1 change: 1 addition & 0 deletions be/src/exprs/agg/factory/aggregate_factory.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@
#include "exprs/agg/hll_union_count.h"
#include "exprs/agg/intersect_count.h"
#include "exprs/agg/mann_whitney.h"
#include "exprs/agg/map_agg.h"
#include "exprs/agg/maxmin.h"
#include "exprs/agg/maxmin_by.h"
#include "exprs/agg/nullable_aggregate.h"
Expand Down
25 changes: 25 additions & 0 deletions be/src/exprs/agg/factory/aggregate_resolver_avg.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -117,13 +117,38 @@ struct ArrayAggDistinctDispatcher {
}
};

struct MapAggDispatcher {
template <LogicalType kt>
void operator()(AggregateFuncResolver* resolver) {
if constexpr (lt_is_aggregate<kt>) {
using KeyCppType = RunTimeCppType<kt>;
if constexpr (lt_is_largeint<kt>) {
using MyHashMap = phmap::flat_hash_map<KeyCppType, Datum, Hash128WithSeed<PhmapSeed1>>;
auto func = std::make_shared<MapAggAggregateFunction<kt, MyHashMap>>();
resolver->add_aggregate_mapping_notnull<kt, TYPE_MAP, AggregateFunctionPtr>("map_agg", false, func);
} else if constexpr (lt_is_fixedlength<kt>) {
using MyHashMap = phmap::flat_hash_map<KeyCppType, Datum, StdHash<KeyCppType>>;
auto func = std::make_shared<MapAggAggregateFunction<kt, MyHashMap>>();
resolver->add_aggregate_mapping_notnull<kt, TYPE_MAP, AggregateFunctionPtr>("map_agg", false, func);
} else if constexpr (lt_is_string<kt>) {
using MyHashMap = phmap::flat_hash_map<SliceWithHash, Datum, HashOnSliceWithHash, EqualOnSliceWithHash>;
auto func = std::make_shared<MapAggAggregateFunction<kt, MyHashMap>>();
resolver->add_aggregate_mapping_notnull<kt, TYPE_MAP, AggregateFunctionPtr>("map_agg", false, func);
} else {
throw std::runtime_error("map_agg does not support key type " + type_to_string(kt));
}
}
}
};

void AggregateFuncResolver::register_avg() {
for (auto type : aggregate_types()) {
type_dispatch_all(type, AvgDispatcher(), this);
type_dispatch_all(type, ArrayAggDispatcher(), this);
type_dispatch_all(type, ArrayAggDistinctDispatcher(), this);
type_dispatch_all(type, ArrayUnionAggDispatcher(), this);
type_dispatch_all(type, ArrayUniqueAggDispatcher(), this);
type_dispatch_all(type, MapAggDispatcher(), this);
}
type_dispatch_all(TYPE_JSON, ArrayAggDispatcher(), this);
add_decimal_mapping<TYPE_DECIMAL32, TYPE_DECIMAL128, true>("decimal_avg");
Expand Down
144 changes: 144 additions & 0 deletions be/src/exprs/agg/map_agg.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,144 @@
// Copyright 2021-present StarRocks, Inc. All rights reserved.
//
// Licensed 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
//
// https://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 <fmt/format.h>

#include "column/binary_column.h"
#include "column/column.h"
#include "column/column_helper.h"
#include "column/fixed_length_column.h"
#include "column/hash_set.h"
#include "column/map_column.h"
#include "column/type_traits.h"
#include "exprs/agg/aggregate.h"
#include "exprs/function_context.h"
#include "gutil/casts.h"
#include "util/phmap/phmap.h"
#include "util/time.h"

namespace starrocks {

template <LogicalType KT, typename MyHashMap = std::map<int, Datum>>
struct MapAggAggregateFunctionState : public AggregateFunctionEmptyState {
using KeyColumnType = RunTimeColumnType<KT>;
using KeyType = typename SliceHashSet::key_type;

MyHashMap hash_map;
void update(MemPool* mem_pool, const KeyColumnType& arg_key_column, const Column& arg_value_column, size_t offset,
size_t count) {
for (int i = offset; i < offset + count; i++) {
if constexpr (!lt_is_string<KT>) {
auto key = arg_key_column.get_data()[i];
auto value = arg_value_column.get(i);
hash_map.emplace(key, value);
} else {
auto raw_key = arg_key_column.get_slice(i);
KeyType key(raw_key);
if (!hash_map.contains(key)) {
uint8_t* pos = mem_pool->allocate(key.size);
memcpy(pos, key.data, key.size);
hash_map.emplace(Slice(pos, key.size), arg_value_column.get(i));
}
}
}
}
};

template <LogicalType KT, typename MyHashMap = std::map<int, int>>
class MapAggAggregateFunction final : public AggregateFunctionBatchHelper<MapAggAggregateFunctionState<KT, MyHashMap>,
MapAggAggregateFunction<KT, MyHashMap>> {
public:
using KeyColumnType = RunTimeColumnType<KT>;
using CppType = RunTimeCppType<KT>;

void update(FunctionContext* ctx, const Column** columns, AggDataPtr __restrict state,
size_t row_num) const override {
// Key could not be null.
if ((columns[0]->is_nullable() && columns[0]->is_null(row_num)) || columns[0]->only_null()) {
return;
}
const auto& key_column = down_cast<const KeyColumnType&>(*ColumnHelper::get_data_column(columns[0]));
this->data(state).update(ctx->mem_pool(), key_column, *columns[1], row_num, 1);
}

void merge(FunctionContext* ctx, const Column* column, AggDataPtr __restrict state, size_t row_num) const override {
auto map_column = down_cast<const MapColumn*>(ColumnHelper::get_data_column(column));
auto& offsets = map_column->offsets().get_data();
if (offsets[row_num + 1] > offsets[row_num]) {
this->data(state).update(
ctx->mem_pool(),
*down_cast<KeyColumnType*>(ColumnHelper::get_data_column(map_column->keys_column().get())),
map_column->values(), offsets[row_num], offsets[row_num + 1] - offsets[row_num]);
}
}

void serialize_to_column(FunctionContext* ctx, ConstAggDataPtr __restrict state, Column* to) const override {
auto& state_impl = this->data(state);
auto* map_column = down_cast<MapColumn*>(ColumnHelper::get_data_column(to));

auto elem_size = state_impl.hash_map.size();

auto* key_column = down_cast<KeyColumnType*>(ColumnHelper::get_data_column(map_column->keys_column().get()));
if constexpr (lt_is_string<KT>) {
for (const auto& entry : state_impl.hash_map) {
key_column->append(Slice(entry.first.data, entry.first.size));
map_column->values_column()->append_datum(entry.second);
}
} else {
for (const auto& entry : state_impl.hash_map) {
key_column->append(entry.first);
map_column->values_column()->append_datum(entry.second);
}
}

if (to->is_nullable()) {
down_cast<NullableColumn*>(to)->null_column_data().emplace_back(0);
}
if (map_column->keys_column()->is_nullable()) {
// Key could not be NULL.
auto* nullable_column = down_cast<NullableColumn*>(map_column->keys_column().get());
nullable_column->null_column_data().resize(nullable_column->null_column_data().size() + elem_size);
}

auto& offsets = map_column->offsets_column()->get_data();
offsets.push_back(offsets.back() + elem_size);
}

void finalize_to_column(FunctionContext* ctx, ConstAggDataPtr __restrict state, Column* to) const override {
serialize_to_column(ctx, state, to);
}

void convert_to_serialize_format(FunctionContext* ctx, const Columns& src, size_t chunk_size,
ColumnPtr* dst) const override {
auto* column = down_cast<MapColumn*>(ColumnHelper::get_data_column(dst->get()));
auto key_column = column->keys_column();
auto value_column = column->values_column();
auto& offsets = column->offsets_column()->get_data();
for (size_t i = 0; i < chunk_size; i++) {
if ((src[0]->is_nullable() && src[0]->is_null(i)) || src[0]->only_null()) {
offsets.push_back(offsets.back());
continue;
}
key_column->append(*src[0], i, 1);
value_column->append(*src[1], i, 1);
offsets.push_back(offsets.back() + 1);
}
}

std::string get_name() const override { return "map_agg"; }
};

} // namespace starrocks
28 changes: 28 additions & 0 deletions fe/fe-core/src/main/java/com/starrocks/catalog/FunctionSet.java
Original file line number Diff line number Diff line change
Expand Up @@ -461,6 +461,9 @@ public class FunctionSet {
public static final String MAP_FROM_ARRAYS = "map_from_arrays";
public static final String MAP_KEYS = "map_keys";
public static final String MAP_SIZE = "map_size";

public static final String MAP_AGG = "map_agg";

public static final String TRANSFORM_VALUES = "transform_values";
public static final String TRANSFORM_KEYS = "transform_keys";

Expand Down Expand Up @@ -1093,6 +1096,9 @@ private void initAggregateBuiltins() {
// Percentile
registerBuiltinPercentileAggFunction();

// map_agg
registerBuiltinMapAggFunction();

// HLL_UNION_AGG
addBuiltin(AggregateFunction.createBuiltin(HLL_UNION_AGG,
Lists.newArrayList(Type.HLL), Type.BIGINT, Type.HLL,
Expand Down Expand Up @@ -1323,6 +1329,28 @@ private void registerBuiltinArrayAggDistinctFunction() {
false, false, false));
}

private void registerBuiltinMapAggFunction() {
for (ScalarType keyType : Type.getNumericTypes()) {
addBuiltin(AggregateFunction.createBuiltin(FunctionSet.MAP_AGG,
Lists.newArrayList(keyType, Type.ANY_ELEMENT), Type.ANY_MAP, null,
false, false, false));
}
for (ScalarType keyType : Type.STRING_TYPES) {
addBuiltin(AggregateFunction.createBuiltin(FunctionSet.MAP_AGG,
Lists.newArrayList(keyType, Type.ANY_ELEMENT), Type.ANY_MAP, null,
false, false, false));
}

for (ScalarType keyType : Type.DATE_TYPES) {
addBuiltin(AggregateFunction.createBuiltin(FunctionSet.MAP_AGG,
Lists.newArrayList(keyType, Type.ANY_ELEMENT), Type.ANY_MAP, null,
false, false, false));
}
addBuiltin(AggregateFunction.createBuiltin(FunctionSet.MAP_AGG,
Lists.newArrayList(Type.TIME, Type.ANY_ELEMENT), Type.ANY_MAP, null,
false, false, false));
}

private void registerBuiltinArrayUniqueAggFunction() {
// array_unique_agg mapping array_agg_distinct while array as input.
for (ScalarType type : Type.getNumericTypes()) {
Expand Down
Loading

0 comments on commit 644a6f9

Please sign in to comment.