be/src/exec/common/distinct_agg_utils.h

// 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 <variant>
#include <vector>

#include "core/arena.h"
#include "core/types.h"
#include "exec/common/hash_table/hash_map_context.h"
#include "exec/common/hash_table/hash_map_util.h"
#include "exec/common/hash_table/ph_hash_map.h"
#include "exec/common/hash_table/ph_hash_set.h"
#include "exec/common/hash_table/string_hash_map.h"

namespace doris {

template <typename T>
struct DistinctHashSetType {
    using HashSet = PHHashSet<T, HashCRC32<T>>;
};

template <>
struct DistinctHashSetType<UInt8> {
    using HashSet = SmallFixedSizeHashSet<UInt8>;
};

template <>
struct DistinctHashSetType<Int8> {
    using HashSet = SmallFixedSizeHashSet<Int8>;
};

template <typename T>
struct DistinctPhase2HashSetType {
    using HashSet = PHHashSet<T, HashMixWrapper<T>>;
};

template <>
struct DistinctPhase2HashSetType<UInt8> {
    using HashSet = SmallFixedSizeHashSet<UInt8>;
};

template <>
struct DistinctPhase2HashSetType<Int8> {
    using HashSet = SmallFixedSizeHashSet<Int8>;
};

template <typename T>
using DistinctData = typename DistinctHashSetType<T>::HashSet;

template <typename T>
using DistinctDataPhase2 = typename DistinctPhase2HashSetType<T>::HashSet;

using DistinctDataWithStringKey = PHHashSet<StringRef>;

// todo: Need to implement StringHashSet like StringHashMap
using DistinctDataWithShortStringKey = PHHashSet<StringRef>;

using DistinctMethodVariants = std::variant<
        std::monostate, MethodSerialized<DistinctDataWithStringKey>,
        MethodOneNumber<UInt8, DistinctData<UInt8>>, MethodOneNumber<UInt16, DistinctData<UInt16>>,
        MethodOneNumber<UInt32, DistinctData<UInt32>>,
        MethodOneNumber<UInt64, DistinctData<UInt64>>,
        MethodStringNoCache<DistinctDataWithShortStringKey>,
        MethodOneNumber<UInt128, DistinctData<UInt128>>,
        MethodOneNumber<UInt256, DistinctData<UInt256>>,
        MethodOneNumber<UInt32, DistinctDataPhase2<UInt32>>,
        MethodOneNumber<UInt64, DistinctDataPhase2<UInt64>>,
        MethodSingleNullableColumn<MethodOneNumber<UInt8, DataWithNullKey<DistinctData<UInt8>>>>,
        MethodSingleNullableColumn<MethodOneNumber<UInt16, DataWithNullKey<DistinctData<UInt16>>>>,
        MethodSingleNullableColumn<MethodOneNumber<UInt32, DataWithNullKey<DistinctData<UInt32>>>>,
        MethodSingleNullableColumn<MethodOneNumber<UInt64, DataWithNullKey<DistinctData<UInt64>>>>,
        MethodSingleNullableColumn<
                MethodOneNumber<UInt32, DataWithNullKey<DistinctDataPhase2<UInt32>>>>,
        MethodSingleNullableColumn<
                MethodOneNumber<UInt64, DataWithNullKey<DistinctDataPhase2<UInt64>>>>,
        MethodSingleNullableColumn<
                MethodOneNumber<UInt128, DataWithNullKey<DistinctData<UInt128>>>>,
        MethodSingleNullableColumn<
                MethodOneNumber<UInt256, DataWithNullKey<DistinctData<UInt256>>>>,
        MethodSingleNullableColumn<
                MethodStringNoCache<DataWithNullKey<DistinctDataWithShortStringKey>>>,
        MethodKeysFixed<DistinctData<UInt64>>, MethodKeysFixed<DistinctData<UInt72>>,
        MethodKeysFixed<DistinctData<UInt96>>, MethodKeysFixed<DistinctData<UInt104>>,
        MethodKeysFixed<DistinctData<UInt128>>, MethodKeysFixed<DistinctData<UInt136>>,
        MethodKeysFixed<DistinctData<UInt256>>>;

struct DistinctDataVariants
        : public DataVariants<DistinctMethodVariants, MethodSingleNullableColumn, MethodOneNumber,
                              DataWithNullKey> {
    void init(const std::vector<DataTypePtr>& data_types, HashKeyType type) {
        bool nullable = data_types.size() == 1 && data_types[0]->is_nullable();
        switch (type) {
        case HashKeyType::serialized:
            method_variant.emplace<MethodSerialized<DistinctDataWithStringKey>>();
            break;
        case HashKeyType::int8_key:
            emplace_single<UInt8, DistinctData<UInt8>>(nullable);
            break;
        case HashKeyType::int16_key:
            emplace_single<UInt16, DistinctData<UInt16>>(nullable);
            break;
        case HashKeyType::int32_key:
            emplace_single<UInt32, DistinctData<UInt32>>(nullable);
            break;
        case HashKeyType::int32_key_phase2:
            emplace_single<UInt32, DistinctDataPhase2<UInt32>>(nullable);
            break;
        case HashKeyType::int64_key:
            emplace_single<UInt64, DistinctData<UInt64>>(nullable);
            break;
        case HashKeyType::int64_key_phase2:
            emplace_single<UInt64, DistinctDataPhase2<UInt64>>(nullable);
            break;
        case HashKeyType::int128_key:
            emplace_single<UInt128, DistinctData<UInt128>>(nullable);
            break;
        case HashKeyType::int256_key:
            emplace_single<UInt256, DistinctData<UInt256>>(nullable);
            break;
        case HashKeyType::string_key:
            if (nullable) {
                method_variant.emplace<MethodSingleNullableColumn<
                        MethodStringNoCache<DataWithNullKey<DistinctDataWithShortStringKey>>>>();
            } else {
                method_variant.emplace<MethodStringNoCache<DistinctDataWithShortStringKey>>();
            }
            break;
        case HashKeyType::fixed64:
            method_variant.emplace<MethodKeysFixed<DistinctData<UInt64>>>(
                    get_key_sizes(data_types));
            break;
        case HashKeyType::fixed72:
            method_variant.emplace<MethodKeysFixed<DistinctData<UInt72>>>(
                    get_key_sizes(data_types));
            break;
        case HashKeyType::fixed96:
            method_variant.emplace<MethodKeysFixed<DistinctData<UInt96>>>(
                    get_key_sizes(data_types));
            break;
        case HashKeyType::fixed104:
            method_variant.emplace<MethodKeysFixed<DistinctData<UInt104>>>(
                    get_key_sizes(data_types));
            break;
        case HashKeyType::fixed128:
            method_variant.emplace<MethodKeysFixed<DistinctData<UInt128>>>(
                    get_key_sizes(data_types));
            break;
        case HashKeyType::fixed136:
            method_variant.emplace<MethodKeysFixed<DistinctData<UInt136>>>(
                    get_key_sizes(data_types));
            break;
        case HashKeyType::fixed256:
            method_variant.emplace<MethodKeysFixed<DistinctData<UInt256>>>(
                    get_key_sizes(data_types));
            break;
        default:
            throw Exception(ErrorCode::INTERNAL_ERROR,
                            "AggregatedDataVariants meet invalid key type, type={}", type);
        }
    }
};

} // namespace doris