code/hyperloglog/hyperloglog_optimized_udf.cpp - master - AMaftuna/Cardinality-Estimation-Algorithm

code/hyperloglog/hyperloglog_optimized_udf.cpp
360 строк13 KB
Abdulloeva Maftuna
hll
22 май 2025, 01:03
22 май 2025, 01:03e492bec
#include "library/cpp/hyperloglog/hyperloglog_optimized.h"

#include <yql/essentials/public/udf/udf_helpers.h>

using namespace NKikimr;
using namespace NUdf;

namespace {

    class TOptimizedHyperLogLog {
    public:
        explicit TOptimizedHyperLogLog(unsigned precision)
            : hll(precision)
        { }

        void Update(ui64 hash) {
            hll.Add(std::to_string(hash));
        }

        void Merge(const TOptimizedHyperLogLog& rh) {
            hll.Merge(rh.hll);
        }

        void Save(IOutputStream& out) const {
            hll.Serialize(out);
        }

        static TOptimizedHyperLogLog Load(IInputStream& in, unsigned precision) {
            auto full = ImprovedHyperLogLog::Deserialize(in);
            if (full.GetPrecision() != precision) {
                throw std::runtime_error("Precision mismatch in HLL::Load");
            }
            TOptimizedHyperLogLog result(precision);
            result.hll = std::move(full);
            return result;
        }

        ui64 Estimate() const {
            return static_cast<ui64>(hll.Estimate());
        }

        unsigned GetPrecision() const {
            return hll.GetPrecision();
        }

    private:
        ImprovedHyperLogLog hll;
    };

    extern const char OptimizedHyperLogLogResourceName[] = "OptimizedHyperLogLog.State";
    using TOptimizedHLLResource = TBoxedResource<TOptimizedHyperLogLog, OptimizedHyperLogLogResourceName>;


    class TOptimizedHyperLogLog_Create: public TBoxedValue {
    public:
        TOptimizedHyperLogLog_Create(TSourcePosition pos)
            : Pos_(pos)
        {}

        static const TStringRef& Name() {
            static auto nameRef = TStringRef::Of("CreateOptimized");
            return nameRef;
        }

        TUnboxedValue Run(const IValueBuilder*, const TUnboxedValuePod* args) const override {
            auto hll = new TOptimizedHLLResource(TOptimizedHyperLogLog(args[1].Get<ui32>()));
            hll->Get()->Update(args[0].Get<ui64>());
            return TUnboxedValuePod(hll);
        }

        static bool DeclareSignature(const TStringRef& name, TType*, IFunctionTypeInfoBuilder& builder, bool typesOnly) {
            if (Name() == name) {
                builder.SimpleSignature<TResource<OptimizedHyperLogLogResourceName>(ui64, ui32)>();
                if (!typesOnly) {
                    builder.Implementation(new TOptimizedHyperLogLog_Create(builder.GetSourcePosition()));
                }
                return true;
            }
            return false;
        }

    private:
        TSourcePosition Pos_;
    };


    class TOptimizedHyperLogLog_AddValue: public TBoxedValue {
    public:
        TOptimizedHyperLogLog_AddValue(TSourcePosition pos)
            : Pos_(pos)
        {}

        static const TStringRef& Name() {
            static auto nameRef = TStringRef::Of("AddValueOptimized");
            return nameRef;
        }

        TUnboxedValue Run(const IValueBuilder*, const TUnboxedValuePod* args) const override {
            auto resource = static_cast<TOptimizedHLLResource*>(args[0].AsBoxed().Get());
            resource->Get()->Update(args[1].Get<ui64>());
            return TUnboxedValuePod(args[0]);
        }

        static bool DeclareSignature(const TStringRef& name, TType*, IFunctionTypeInfoBuilder& builder, bool typesOnly) {
            if (Name() == name) {
                builder.SimpleSignature<TResource<OptimizedHyperLogLogResourceName>(TResource<OptimizedHyperLogLogResourceName>, ui64)>();
                if (!typesOnly) {
                    builder.Implementation(new TOptimizedHyperLogLog_AddValue(builder.GetSourcePosition()));
                }
                return true;
            }
            return false;
        }

    private:
        TSourcePosition Pos_;
    };

    class TOptimizedHyperLogLog_GetResult: public TBoxedValue {
    public:
        TOptimizedHyperLogLog_GetResult(TSourcePosition pos)
            : Pos_(pos)
        {}

        static const TStringRef& Name() {
            static auto nameRef = TStringRef::Of("GetResultOptimized");
            return nameRef;
        }

        TUnboxedValue Run(const IValueBuilder*, const TUnboxedValuePod* args) const override {
            auto hll = static_cast<TOptimizedHLLResource*>(args[0].AsBoxed().Get())->Get();
            return TUnboxedValuePod(hll->Estimate());
        }

        static bool DeclareSignature(const TStringRef& name, TType*, IFunctionTypeInfoBuilder& builder, bool typesOnly) {
            if (Name() == name) {
                builder.SimpleSignature<ui64(TResource<OptimizedHyperLogLogResourceName>)>();
                if (!typesOnly) {
                    builder.Implementation(new TOptimizedHyperLogLog_GetResult(builder.GetSourcePosition()));
                }
                return true;
            }
            return false;
        }

    private:
        TSourcePosition Pos_;
    };

    class TOptimizedHyperLogLog_Distinct : public TBoxedValue {
    public:
        TOptimizedHyperLogLog_Distinct(TSourcePosition pos)
            : Pos_(pos)
        {}

        static const TStringRef& Name() {
            static auto nameRef = TStringRef::Of("DistinctOptimized");
            return nameRef;
        }

        TUnboxedValue Run(const IValueBuilder*, const TUnboxedValuePod* args) const override {
            auto hll = TOptimizedHyperLogLog(args[1].Get<ui32>());
            hll.Update(args[0].Get<ui64>());
            return TUnboxedValuePod(hll.Estimate());
        }

        static bool DeclareSignature(const TStringRef& name, TType*, IFunctionTypeInfoBuilder& builder, bool typesOnly) {
            if (Name() == name) {
                builder.SimpleSignature<double(ui64, float)>();
                if (!typesOnly) {
                    builder.Implementation(new TOptimizedHyperLogLog_Distinct(builder.GetSourcePosition()));
                }
                return true;
            }
            return false;
        }

    private:
        TSourcePosition Pos_;
    };

    class TOptimizedHyperLogLog_Accum : public TBoxedValue {
    public:
        TOptimizedHyperLogLog_Accum(TSourcePosition pos)
            : Pos_(pos)
        {}

        static const TStringRef& Name() {
            static auto nameRef = TStringRef::Of("AccumOptimized");
            return nameRef;
        }

        TUnboxedValue Run(const IValueBuilder*, const TUnboxedValuePod* args) const override {
            auto resource = new TOptimizedHLLResource(TOptimizedHyperLogLog(args[1].Get<ui32>()));
            resource->Get()->Update(args[0].Get<ui64>());
            return TUnboxedValuePod(resource);
        }

        static bool DeclareSignature(const TStringRef& name, TType*, IFunctionTypeInfoBuilder& builder, bool typesOnly) {
            if (Name() == name) {
                builder.SimpleSignature<TResource<OptimizedHyperLogLogResourceName>(ui64, float)>();
                if (!typesOnly) {
                    builder.Implementation(new TOptimizedHyperLogLog_Accum(builder.GetSourcePosition()));
                }
                return true;
            }
            return false;
        }

    private:
        TSourcePosition Pos_;
    };

    class TOptimizedHyperLogLog_Sum : public TBoxedValue {
    public:
        TOptimizedHyperLogLog_Sum(TSourcePosition pos)
            : Pos_(pos)
        {}

        static const TStringRef& Name() {
            static auto nameRef = TStringRef::Of("SumOptimized");
            return nameRef;
        }

        TUnboxedValue Run(const IValueBuilder*, const TUnboxedValuePod* args) const override {
            auto resource = static_cast<TOptimizedHLLResource*>(args[0].AsBoxed().Get())->Get();
            return TUnboxedValuePod(resource->Estimate());
        }

        static bool DeclareSignature(const TStringRef& name, TType*, IFunctionTypeInfoBuilder& builder, bool typesOnly) {
            if (Name() == name) {
                builder.SimpleSignature<ui64(TResource<OptimizedHyperLogLogResourceName>)>();
                if (!typesOnly) {
                    builder.Implementation(new TOptimizedHyperLogLog_Sum(builder.GetSourcePosition()));
                }
                return true;
            }
            return false;
        }

    private:
        TSourcePosition Pos_;
    };

    class TOptimizedHyperLogLog_Merge : public TBoxedValue {
    public:
        TOptimizedHyperLogLog_Merge(TSourcePosition pos)
            : Pos_(pos)
        {}

        static const TStringRef& Name() {
            static auto nameRef = TStringRef::Of("MergeOptimized");
            return nameRef;
        }

        TUnboxedValue Run(const IValueBuilder*, const TUnboxedValuePod* args) const override {
            auto left = static_cast<TOptimizedHLLResource*>(args[0].AsBoxed().Get())->Get();
            auto right = static_cast<TOptimizedHLLResource*>(args[1].AsBoxed().Get())->Get();
            right->Merge(*left);
            return TUnboxedValuePod(args[1]);
        }

        static bool DeclareSignature(const TStringRef& name, TType*, IFunctionTypeInfoBuilder& builder, bool typesOnly) {
            if (Name() == name) {
                builder.SimpleSignature<TResource<OptimizedHyperLogLogResourceName>(TResource<OptimizedHyperLogLogResourceName>, TResource<OptimizedHyperLogLogResourceName>)>();
                if (!typesOnly) {
                    builder.Implementation(new TOptimizedHyperLogLog_Merge(builder.GetSourcePosition()));
                }
                return true;
            }
            return false;
        }

    private:
        TSourcePosition Pos_;
    };

    class TOptimizedHyperLogLog_Serialize : public TBoxedValue {
    public:
        TOptimizedHyperLogLog_Serialize(TSourcePosition pos)
            : Pos_(pos)
        {}

        static const TStringRef& Name() {
            static auto nameRef = TStringRef::Of("SerializeOptimized");
            return nameRef;
        }

        TUnboxedValue Run(const IValueBuilder* builder, const TUnboxedValuePod* args) const override {
            TStringStream out;
            auto resource = static_cast<TOptimizedHLLResource*>(args[0].AsBoxed().Get())->Get();
            ui32 precision = resource->GetPrecision();
            out.Write(&precision, sizeof(precision));
            resource->Save(out);
            return builder->NewString(out.Str());
        }

        static bool DeclareSignature(const TStringRef& name, TType*, IFunctionTypeInfoBuilder& builder, bool typesOnly) {
            if (Name() == name) {
                builder.SimpleSignature<char*(TResource<OptimizedHyperLogLogResourceName>)>();
                if (!typesOnly) {
                    builder.Implementation(new TOptimizedHyperLogLog_Serialize(builder.GetSourcePosition()));
                }
                return true;
            }
            return false;
        }

    private:
        TSourcePosition Pos_;
    };

    class TOptimizedHyperLogLog_Deserialize : public TBoxedValue {
    public:
        TOptimizedHyperLogLog_Deserialize(TSourcePosition pos)
            : Pos_(pos)
        {}

        static const TStringRef& Name() {
            static auto nameRef = TStringRef::Of("DeserializeOptimized");
            return nameRef;
        }

        TUnboxedValue Run(const IValueBuilder*, const TUnboxedValuePod* args) const override {
            const TString serialized(args[0].AsStringRef());
            TStringInput input(serialized);
            ui32 precision;
            input.Read(&precision, sizeof(precision));
            auto resource = new TOptimizedHLLResource(TOptimizedHyperLogLog::Load(input, precision));
            return TUnboxedValuePod(resource);
        }

        static bool DeclareSignature(const TStringRef& name, TType*, IFunctionTypeInfoBuilder& builder, bool typesOnly) {
            if (Name() == name) {
                builder.SimpleSignature<TResource<OptimizedHyperLogLogResourceName>(char*)>();
                if (!typesOnly) {
                    builder.Implementation(new TOptimizedHyperLogLog_Deserialize(builder.GetSourcePosition()));
                }
                return true;
            }
            return false;
        }

    private:
        TSourcePosition Pos_;
    };

    SIMPLE_MODULE(TOptimizedHyperLogLogModule,
                  TOptimizedHyperLogLog_Create,
                  TOptimizedHyperLogLog_AddValue,
                  TOptimizedHyperLogLog_GetResult,
                  TOptimizedHyperLogLog_Distinct,
                  TOptimizedHyperLogLog_Accum,
                  TOptimizedHyperLogLog_Sum,
                  TOptimizedHyperLogLog_Merge,
                  TOptimizedHyperLogLog_Serialize,
                  TOptimizedHyperLogLog_Deserialize)
}

REGISTER_MODULES(TOptimizedHyperLogLogModule)