statistics: add count-min sketch UDAF implementation (#28786)

Author: ztlpn

ydb/apps/ydbd/ya.make

@@ -48,6 +48,7 @@ PEERDIR(
     ydb/library/yql/udfs/common/clickhouse/client
     ydb/library/yql/udfs/common/knn
     ydb/library/yql/udfs/common/roaring
+    ydb/library/yql/udfs/statistics_internal
     yql/essentials/parser/pg_wrapper
     yql/essentials/sql/pg
     yql/essentials/udfs/common/compress_base

ydb/library/yql/udfs/statistics_internal/README.md

@@ -0,0 +1,5 @@
+# StatisticsInternal YQL UDF module
+
+This module contains the implementation of UDAFs needed for column statistics calculation.
+
+They are not intended to be invoked by end users.

ydb/library/yql/udfs/statistics_internal/all_agg_funcs.cpp

@@ -0,0 +1 @@
+#include <ydb/library/yql/udfs/statistics_internal/all_agg_funcs.h>

ydb/library/yql/udfs/statistics_internal/all_agg_funcs.h

@@ -0,0 +1,42 @@
+#pragma once
+
+#include <ydb/core/scheme/scheme_tablecell.h>
+#include <yql/essentials/core/minsketch/count_min_sketch.h>
+
+namespace NKikimr::NStat::NAggFuncs {
+
+class TCMSAggFunc {
+public:
+    static constexpr std::string_view GetName() { return "CountMinSketch"; }
+
+    using TState = std::unique_ptr<NKikimr::TCountMinSketch>;
+
+    static constexpr size_t ParamsCount = 2;
+
+    static TState CreateState(const std::array<ui64, ParamsCount>& params) {
+        return TState(NKikimr::TCountMinSketch::Create(params[0], params[1]));
+    }
+    static auto CreateStateUpdater() {
+        return [](TState& state, const TCell& cell) {
+            state->Count(cell.Data(), cell.Size());
+        };
+    }
+    static void MergeStates(const TState& left, TState& right) {
+        *right += *left;
+    }
+    static TString SerializeState(const TState& state) {
+        return TString(state->AsStringBuf());
+    }
+    static TState DeserializeState(const char* data, size_t size) {
+        return TState(NKikimr::TCountMinSketch::FromString(data, size));
+    }
+    static TString FinalizeState(const TState& state) {
+        return SerializeState(state);
+    }
+};
+
+using TAllAggFuncsList = TTypeList<
+    TCMSAggFunc
+>;
+
+} // NKikimr::NStat::NAggFuncs

ydb/library/yql/udfs/statistics_internal/test/canondata/result.json

@@ -0,0 +1,7 @@
+{
+    "test.test[count_min_sketch]": [
+        {
+            "uri": "file://test.test_count_min_sketch_/results.txt"
+        }
+    ]
+}

ydb/library/yql/udfs/statistics_internal/test/canondata/test.test_count_min_sketch_/results.txt

@@ -0,0 +1,46 @@
+[
+    {
+        "Write" = [
+            {
+                "Type" = [
+                    "ListType";
+                    [
+                        "StructType";
+                        [
+                            [
+                                "column0";
+                                [
+                                    "OptionalType";
+                                    [
+                                        "DataType";
+                                        "String"
+                                    ]
+                                ]
+                            ];
+                            [
+                                "column1";
+                                [
+                                    "OptionalType";
+                                    [
+                                        "DataType";
+                                        "String"
+                                    ]
+                                ]
+                            ]
+                        ]
+                    ]
+                ];
+                "Data" = [
+                    [
+                        [
+                            "\2\0\0\0\0\0\0\0\2\0\0\0\0\0\0\0\3\0\0\0\0\0\0\0\2\0\0\0\1\0\0\0\1\0\0\0\2\0\0\0"
+                        ];
+                        [
+                            "\3\0\0\0\0\0\0\0\3\0\0\0\0\0\0\0\3\0\0\0\0\0\0\0\2\0\0\0\0\0\0\0\1\0\0\0\1\0\0\0\0\0\0\0\2\0\0\0\1\0\0\0\1\0\0\0\1\0\0\0"
+                        ]
+                    ]
+                ]
+            }
+        ]
+    }
+]

ydb/library/yql/udfs/statistics_internal/test/cases/count_min_sketch.sql

@@ -0,0 +1,18 @@
+$get_factory = ($width, $depth) -> { return AggregationFactory(
+        "UDAF",
+        ($item, $parent) -> { return Udf(StatisticsInternal::CountMinSketchCreate, $parent as Depends)($item, $width, $depth) },
+        ($state, $item, $parent) -> { return Udf(StatisticsInternal::CountMinSketchAddValue, $parent as Depends)($state, $item) },
+        StatisticsInternal::CountMinSketchMerge,
+        StatisticsInternal::CountMinSketchFinalize,
+        StatisticsInternal::CountMinSketchSerialize,
+        StatisticsInternal::CountMinSketchDeserialize,
+    )
+};
+
+$t1 = [
+    <|key: 1, value: "v1"|>,
+    <|key: 2, value: "v2"|>,
+    <|key: 3, value: "v3"|>
+];
+
+select AGGREGATE_BY(key, $get_factory(2, 2)), AGGREGATE_BY(value, $get_factory(3, 3)) from AS_TABLE($t1);

ydb/library/yql/udfs/statistics_internal/test/ya.make

@@ -0,0 +1,7 @@
+YQL_UDF_TEST()
+
+DEPENDS(ydb/library/yql/udfs/statistics_internal)
+
+SIZE(MEDIUM)
+
+END()