support DataSketches quantile sketche

Signed-off-by: chenminghua8 <cmptmn@126.com>

be/src/exprs/agg/ds_agg.h

@@ -15,6 +15,7 @@
 #pragma once
 
 #include "exprs/agg/ds_hll_state.h"
+#include "exprs/agg/ds_quantile_state.h"
 #include "exprs/agg/ds_state.h"
 
 namespace starrocks {

be/src/exprs/agg/ds_quantile_state.h

@@ -0,0 +1,176 @@
+// 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 "column/array_column.h"
+#include "column/column_helper.h"
+#include "exprs/agg/ds_state.h"
+#include "types/ds_quantile_sketch.cpp"
+
+namespace starrocks {
+template<LogicalType LT>
+struct DSSketchState<LT, QUANTILE> {
+    using CppType = RunTimeCppType<LT>;
+    using ColumnType = RunTimeColumnType<LT>;
+    using SketchWarapperType = DataSketchesQuantile<CppType>;
+    uint32_t ranks_size;
+    std::unique_ptr<double[]> ranks = nullptr;
+    std::unique_ptr <SketchWarapperType> ds_sketch_wrapper = nullptr;
+    int64_t memory_usage = 0;
+
+    void init(FunctionContext *ctx) {
+        DatumArray datum_array;
+        uint16_t k;
+        std::tie(k, datum_array) = _parse_sketch_args(ctx);
+        if (datum_array.size() < 1) {
+            ranks_size = 1;
+            ranks = std::make_unique<double[]>(ranks_size);
+            *ranks.get() = 0.5;
+        } else {
+            ranks_size = datum_array.size();
+            ranks = std::make_unique<double[]>(ranks_size);
+            double *ranks_prt = ranks.get();
+            for (Datum rank: datum_array) {
+                *ranks_prt = rank.get_double();
+                ranks_prt++;
+            }
+        }
+        if (ranks_size == 0) {
+            ranks_size = 0;
+        }
+        ds_sketch_wrapper = std::make_unique<SketchWarapperType>(k, &memory_usage);
+    }
+
+    bool is_inited() const { return ds_sketch_wrapper != nullptr; }
+
+    void update(const Column *data_column, size_t row_num) const {
+        const ColumnType *column = down_cast<const ColumnType *>(data_column);
+        const auto &values = column->get_data();
+        ds_sketch_wrapper->update(values[row_num]);
+    }
+
+    void
+    update_batch_single_state_with_frame(const Column *data_column, int64_t frame_start, int64_t frame_end) const {
+        const ColumnType *column = down_cast<const ColumnType *>(data_column);
+        const auto &values = column->get_data();
+        for (size_t i = frame_start; i < frame_end; ++i) {
+            ds_sketch_wrapper->update(values[i]);
+        }
+    }
+
+    void merge(const BinaryColumn *sketch_data_column, size_t row_num) {
+        DSSketchState<LT, QUANTILE> other_state;
+        other_state.deserialize(sketch_data_column->get(row_num).get_slice(), &memory_usage);
+        if (UNLIKELY(!is_inited())) {
+            ranks_size = other_state.ranks_size;
+            ranks = std::make_unique<double[]>(ranks_size);
+            double *ranks_prt = ranks.get();
+            for (int i = 0; i < ranks_size; i++) {
+                *ranks_prt = other_state.ranks.get()[i];
+                ranks_prt++;
+            }
+            ds_sketch_wrapper =
+                    std::make_unique<SketchWarapperType>(other_state.ds_sketch_wrapper->get_k(), &memory_usage);
+            }
+        ds_sketch_wrapper->merge(*other_state.ds_sketch_wrapper);
+    }
+
+    size_t serialize(uint8_t *dst) const {
+        size_t offset = 0;
+        memcpy(dst + offset, &ranks_size, sizeof(ranks_size));
+        offset = offset + sizeof(uint32_t);
+        memcpy(dst + offset, ranks.get(), ranks_size * sizeof(double));
+        offset = offset + ranks_size * sizeof(double);
+        size_t ser_sketch_size = ds_sketch_wrapper->serialize(dst + offset);
+        return offset + ser_sketch_size;
+    }
+
+    size_t serialize_size() const {
+        return sizeof(uint32_t) + ranks_size * sizeof(double) + ds_sketch_wrapper->serialize_size();
+    }
+
+    void deserialize(const Slice &slice, int64_t *memory_usage) {
+        uint8_t *ptr = (uint8_t *) slice.get_data();
+        size_t offset = 0;
+        memcpy(&ranks_size, ptr + offset, sizeof(uint32_t));
+        if (ranks_size == 0) {
+            ranks_size = 0;
+        }
+        offset = offset + sizeof(uint32_t);
+        ranks = std::make_unique<double[]>(ranks_size);
+        memcpy(ranks.get(), ptr + offset, ranks_size * sizeof(double));
+        offset = offset + ranks_size * sizeof(double);
+        const Slice sketch_data_slice = Slice(slice.get_data() + offset, slice.size - offset);
+        ds_sketch_wrapper = std::make_unique<SketchWarapperType>(sketch_data_slice, memory_usage);
+    }
+
+    void get_values(Column *dst, size_t start, size_t end) const {
+        auto *array_column = down_cast<ArrayColumn *>(dst);
+        auto &offset_column = array_column->offsets_column();
+        auto &elements_column = array_column->elements_column();
+        auto *nullable_column = down_cast<NullableColumn *>(elements_column.get());
+        auto *result_column = down_cast<ColumnType *>(nullable_column->data_column().get());
+
+        std::vector <CppType> result;
+        if (LIKELY(ds_sketch_wrapper != nullptr)) {
+            result = ds_sketch_wrapper->get_quantiles(ranks.get(), ranks_size);
+        }
+
+        uint32_t index = 0;
+        for (size_t row = start; row < end; row++) {
+            for (CppType result_data: result) {
+                result_column->append(result_data);
+                nullable_column->null_column()->append(0);
+                index++;
+            }
+            offset_column->append(index);
+        }
+    }
+
+    static std::string getFunName() { return "ds_quantile"; }
+
+private:
+    // parse k and rank_arr from args
+    static std::tuple <uint16_t, DatumArray> _parse_sketch_args(FunctionContext *ctx) {
+        uint16_t k = DEFAULT_QUANTILE_K;
+        if (ctx->get_num_constant_columns() > 1) {
+            if (ctx->get_num_constant_columns() > 2) {
+                k = ColumnHelper::get_const_value<TYPE_INT>(ctx->get_constant_column(2));
+                if (k <= 1) {
+                    k = DEFAULT_QUANTILE_K;
+                }
+                int i = 1;
+                while ((1 << i) < k) {
+                    i += 1;
+                }
+                k = 1 << i;
+            }
+            Column *ranks_column = ColumnHelper::get_data_column(ctx->get_constant_column(1).get());
+            if (ranks_column->is_array()) {
+                DatumArray rank_arr = ranks_column->get(0).get_array();
+                return {k, rank_arr};
+            } else {
+                DatumArray rank_arr;
+                double rank_value = ranks_column->get(0).get_double();
+                rank_arr.push_back(rank_value);
+                return {k, rank_arr};
+            }
+        }
+        DatumArray rank_arr;
+        return {k, rank_arr};
+    }
+};
+
+} // namespace starrocks

be/src/types/CMakeLists.txt

@@ -18,6 +18,7 @@ add_library(Types STATIC
     array_type_info.cpp
     bitmap_value.cpp
     date_value.cpp
+    ds_quantile_sketch.cpp
     hll.cpp
     hll_sketch.cpp
     large_int_value.cpp

be/src/types/ds_quantile_sketch.cpp

@@ -0,0 +1,97 @@
+// 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.
+
+#include "types/ds_quantile_sketch.h"
+
+namespace starrocks {
+
+    template <typename T>
+    void DataSketchesQuantile<T>::update(T value) { _sketch->update(value); }
+
+    template <typename T>
+    void DataSketchesQuantile<T>::merge(const DataSketchesQuantile<T>& other) {
+        if (other._sketch == nullptr) {
+            return;
+        }
+        if (UNLIKELY(_sketch == nullptr)) {
+            _sketch = std::make_unique<quantile_sketch_type>(other._sketch->get_k(), std::less<T>(),
+                                                             alloc_type(_memory_usage));
+        }
+        _sketch.get()->merge(*other._sketch);
+    }
+
+    template <typename T>
+    uint64_t DataSketchesQuantile<T>::serialize_size() const {
+        if (_sketch == nullptr) {
+            return 0;
+        }
+        return _sketch->get_serialized_size_bytes();
+    }
+
+    template <typename T>
+    bool DataSketchesQuantile<T>::deserialize(const Slice& slice) {
+        DCHECK(_sketch == nullptr);
+
+        if (!is_valid(slice)) {
+            return false;
+        }
+        try {
+            _sketch = std::make_unique<quantile_sketch_type>(
+                    quantile_sketch_type::deserialize((uint8_t*)slice.data, slice.size, datasketches::serde<T>(),
+                                                      std::less<T>(), alloc_type(_memory_usage)));
+        } catch (std::logic_error& e) {
+            LOG(WARNING) << "DataSketchesQuantile deserialize error: " << e.what();
+            return false;
+        }
+        return true;
+    }
+
+    template <typename T>
+    std::vector<T> DataSketchesQuantile<T>::get_quantiles(const double* ranks, uint32_t size) const {
+        std::vector<T> result;
+        if (_sketch == nullptr) {
+            return result;
+        }
+        try {
+            std::vector<T, alloc_type> quantiles = _sketch->get_quantiles(ranks, size);
+            for (T quantile : quantiles) {
+                result.push_back(quantile);
+            }
+        } catch (std::logic_error& e) {
+            LOG(WARNING) << "DataSketchesQuantile get_quantiles error: " << e.what();
+            result.clear();
+        }
+        return result;
+    }
+
+    template <typename T>
+    void DataSketchesQuantile<T>::clear() {
+        *_memory_usage = 0;
+        this->_sketch =
+                std::make_unique<quantile_sketch_type>(_sketch->get_k(), std::less<T>(), alloc_type(_memory_usage));
+    }
+
+    template <typename T>
+    std::string DataSketchesQuantile<T>::to_string() const {
+        if (_sketch == nullptr) {
+            return "";
+        }
+        datasketches::string<alloc_type> str = _sketch->to_string();
+        return std::string(str.begin(), str.end());
+    }
+
+} // namespace starrocks

be/src/types/ds_quantile_sketch.h

@@ -0,0 +1,94 @@
+// 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 <datasketches/quantiles_sketch.hpp>
+#include "runtime/memory/counting_allocator.h"
+#include "util/slice.h"
+
+namespace starrocks {
+template <typename T>
+class DataSketchesQuantile {
+public:
+    using alloc_type = STLCountingAllocator<T>;
+    using quantile_sketch_type = datasketches::quantiles_sketch<T, std::less<T>, alloc_type>;
+
+    explicit DataSketchesQuantile(uint16_t k, int64_t* memory_usage) : _memory_usage(memory_usage) {
+        this->_sketch = std::make_unique<quantile_sketch_type>(k, std::less<T>(), alloc_type(_memory_usage));
+    }
+
+    DataSketchesQuantile(const DataSketchesQuantile& other) = delete;
+    DataSketchesQuantile& operator=(const DataSketchesQuantile& other) = delete;
+
+    DataSketchesQuantile(DataSketchesQuantile&& other) noexcept
+            : _memory_usage(std::move(other._memory_usage)), _sketch(std::move(other._sketch)) {}
+    DataSketchesQuantile& operator=(DataSketchesQuantile&& other) noexcept {
+        if (this != &other) {
+            this->_memory_usage = std::move(other._memory_usage);
+            this->_sketch = std::move(other._sketch);
+        }
+        return *this;
+    }
+
+    explicit DataSketchesQuantile(const Slice& src, int64_t* memory_usage) : _memory_usage(memory_usage) {
+        if (!deserialize(src)) {
+            LOG(WARNING) << "Failed to init DataSketchesQuantile from slice, will be reset to 0.";
+        }
+    }
+
+    ~DataSketchesQuantile() = default;
+
+    uint16_t get_k() const { return _sketch->get_k(); }
+
+    void update(T value);
+
+    void merge(const DataSketchesQuantile& other);
+
+    int64_t mem_usage() const { return _memory_usage == nullptr ? 0L : *_memory_usage; }
+
+    size_t serialize(uint8_t* dst) const {
+        if (_sketch == nullptr) {
+            return 0;
+        }
+        auto serialize_compact = _sketch->serialize();
+        std::copy(serialize_compact.begin(), serialize_compact.end(), dst);
+        return _sketch->get_serialized_size_bytes();
+    }
+
+    uint64_t serialize_size() const;
+
+    bool deserialize(const Slice& slice);
+
+    std::vector<T> get_quantiles(const double* ranks, uint32_t size) const;
+
+    static bool is_valid(const Slice& slice) {
+        if (slice.size < 1) {
+            return false;
+        }
+        return true;
+    }
+
+    void clear();
+
+    std::string to_string() const;
+
+private:
+    int64_t* _memory_usage;
+    mutable std::unique_ptr<quantile_sketch_type> _sketch = nullptr;
+};
+} // namespace starrocks