From 9dac77266c086187655e8aaf57b704c383521f61 Mon Sep 17 00:00:00 2001 From: "liangjie.liang" Date: Wed, 29 Jul 2026 17:02:05 +0800 Subject: [PATCH] feat(aggregate): add Java-compatible field aggregators Co-Authored-By: Claude Opus 5 --- cmake_modules/DefineOptions.cmake | 6 + cmake_modules/FindDataSketchesAlt.cmake | 39 ++ cmake_modules/ThirdpartyToolchain.cmake | 36 ++ include/paimon/defs.h | 8 + src/paimon/CMakeLists.txt | 16 + src/paimon/common/data/generic_array.cpp | 164 ++++++++ src/paimon/common/data/generic_array.h | 79 ++++ src/paimon/common/data/generic_map.cpp | 42 ++ src/paimon/common/data/generic_map.h | 44 +++ src/paimon/common/defs.cpp | 4 + src/paimon/core/core_options.cpp | 55 +++ src/paimon/core/core_options.h | 30 ++ .../aggregate/aggregate_merge_function.cpp | 3 +- .../aggregate/field_aggregate_utils.cpp | 237 +++++++++++ .../compact/aggregate/field_aggregate_utils.h | 62 +++ .../compact/aggregate/field_aggregator.h | 21 + .../aggregate/field_aggregator_factory.h | 23 ++ .../field_aggregator_factory_test.cpp | 29 ++ .../compact/aggregate/field_collect_agg.cpp | 175 +++++++++ .../compact/aggregate/field_collect_agg.h | 67 ++++ .../aggregate/field_collect_agg_test.cpp | 92 +++++ .../aggregate/field_ignore_retract_agg.h | 5 + .../field_ignore_retract_agg_test.cpp | 29 ++ .../compact/aggregate/field_merge_map_agg.cpp | 157 ++++++++ .../compact/aggregate/field_merge_map_agg.h | 60 +++ .../aggregate/field_merge_map_agg_test.cpp | 87 ++++ .../aggregate/field_nested_update_agg.cpp | 370 ++++++++++++++++++ .../aggregate/field_nested_update_agg.h | 78 ++++ .../field_nested_update_agg_test.cpp | 190 +++++++++ .../aggregate/field_roaring_bitmap64_agg.cpp | 62 +++ .../aggregate/field_roaring_bitmap64_agg.h | 48 +++ .../field_roaring_bitmap64_agg_test.cpp | 62 +++ .../compact/aggregate/field_sketch_agg.cpp | 136 +++++++ .../compact/aggregate/field_sketch_agg.h | 70 ++++ .../aggregate/field_sketch_agg_test.cpp | 186 +++++++++ .../compact/partial_update_merge_function.cpp | 24 +- .../compact/partial_update_merge_function.h | 4 +- third_party/versions.txt | 5 + 38 files changed, 2795 insertions(+), 10 deletions(-) create mode 100644 cmake_modules/FindDataSketchesAlt.cmake create mode 100644 src/paimon/common/data/generic_array.cpp create mode 100644 src/paimon/common/data/generic_array.h create mode 100644 src/paimon/common/data/generic_map.cpp create mode 100644 src/paimon/common/data/generic_map.h create mode 100644 src/paimon/core/mergetree/compact/aggregate/field_aggregate_utils.cpp create mode 100644 src/paimon/core/mergetree/compact/aggregate/field_aggregate_utils.h create mode 100644 src/paimon/core/mergetree/compact/aggregate/field_collect_agg.cpp create mode 100644 src/paimon/core/mergetree/compact/aggregate/field_collect_agg.h create mode 100644 src/paimon/core/mergetree/compact/aggregate/field_collect_agg_test.cpp create mode 100644 src/paimon/core/mergetree/compact/aggregate/field_merge_map_agg.cpp create mode 100644 src/paimon/core/mergetree/compact/aggregate/field_merge_map_agg.h create mode 100644 src/paimon/core/mergetree/compact/aggregate/field_merge_map_agg_test.cpp create mode 100644 src/paimon/core/mergetree/compact/aggregate/field_nested_update_agg.cpp create mode 100644 src/paimon/core/mergetree/compact/aggregate/field_nested_update_agg.h create mode 100644 src/paimon/core/mergetree/compact/aggregate/field_nested_update_agg_test.cpp create mode 100644 src/paimon/core/mergetree/compact/aggregate/field_roaring_bitmap64_agg.cpp create mode 100644 src/paimon/core/mergetree/compact/aggregate/field_roaring_bitmap64_agg.h create mode 100644 src/paimon/core/mergetree/compact/aggregate/field_roaring_bitmap64_agg_test.cpp create mode 100644 src/paimon/core/mergetree/compact/aggregate/field_sketch_agg.cpp create mode 100644 src/paimon/core/mergetree/compact/aggregate/field_sketch_agg.h create mode 100644 src/paimon/core/mergetree/compact/aggregate/field_sketch_agg_test.cpp diff --git a/cmake_modules/DefineOptions.cmake b/cmake_modules/DefineOptions.cmake index 4aacc33da..37b5d4453 100644 --- a/cmake_modules/DefineOptions.cmake +++ b/cmake_modules/DefineOptions.cmake @@ -223,6 +223,12 @@ if("${CMAKE_SOURCE_DIR}" STREQUAL "${CMAKE_CURRENT_SOURCE_DIR}") AUTO BUNDLED SYSTEM) + define_option_string(DataSketches_SOURCE + "Dependency source for DataSketches" + "" + AUTO + BUNDLED + SYSTEM) define_option_string(TBB_SOURCE "Dependency source for TBB" "" diff --git a/cmake_modules/FindDataSketchesAlt.cmake b/cmake_modules/FindDataSketchesAlt.cmake new file mode 100644 index 000000000..ed7f8a96f --- /dev/null +++ b/cmake_modules/FindDataSketchesAlt.cmake @@ -0,0 +1,39 @@ +# Copyright 2026-present Alibaba Inc. +# +# 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 +# +# 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. + +# Headers are located directly: the installed config exports unnamespaced targets like `hll`. + +set(_PAIMON_DATASKETCHES_ROOTS ${DataSketches_ROOT} ${DATASKETCHES_ROOT} + ${PAIMON_PACKAGE_PREFIX}) +list(REMOVE_ITEM _PAIMON_DATASKETCHES_ROOTS "") +if(_PAIMON_DATASKETCHES_ROOTS) + set(_PAIMON_DATASKETCHES_FIND_ARGS HINTS ${_PAIMON_DATASKETCHES_ROOTS} + NO_DEFAULT_PATH) +endif() + +find_path(DATASKETCHES_INCLUDE_DIR + NAMES DataSketches/hll.hpp ${_PAIMON_DATASKETCHES_FIND_ARGS} + PATH_SUFFIXES include) + +include(FindPackageHandleStandardArgs) +find_package_handle_standard_args(DataSketchesAlt REQUIRED_VARS DATASKETCHES_INCLUDE_DIR) + +if(DataSketchesAlt_FOUND AND NOT TARGET DataSketches) + add_library(DataSketches INTERFACE IMPORTED) + target_include_directories(DataSketches SYSTEM + INTERFACE "${DATASKETCHES_INCLUDE_DIR}") +endif() + +unset(_PAIMON_DATASKETCHES_FIND_ARGS) +unset(_PAIMON_DATASKETCHES_ROOTS) diff --git a/cmake_modules/ThirdpartyToolchain.cmake b/cmake_modules/ThirdpartyToolchain.cmake index 872c70b8c..8bf376a26 100644 --- a/cmake_modules/ThirdpartyToolchain.cmake +++ b/cmake_modules/ThirdpartyToolchain.cmake @@ -103,6 +103,19 @@ else() endif() endif() +if(DEFINED ENV{PAIMON_DATASKETCHES_URL}) + set(DATASKETCHES_SOURCE_URL "$ENV{PAIMON_DATASKETCHES_URL}") +else() + if(EXISTS "${THIRDPARTY_DIR}/${PAIMON_DATASKETCHES_PKG_NAME}") + set_urls(DATASKETCHES_SOURCE_URL + "${THIRDPARTY_DIR}/${PAIMON_DATASKETCHES_PKG_NAME}") + else() + set_urls(DATASKETCHES_SOURCE_URL + "${THIRDPARTY_MIRROR_URL}https://github.com/apache/datasketches-cpp/archive/refs/tags/${PAIMON_DATASKETCHES_BUILD_VERSION}.tar.gz" + ) + endif() +endif() + if(DEFINED ENV{PAIMON_FMT_URL}) set(FMT_SOURCE_URL "$ENV{PAIMON_FMT_URL}") else() @@ -600,6 +613,8 @@ macro(paimon_build_dependency DEPENDENCY_NAME) build_fmt() elseif("${DEPENDENCY_NAME}" STREQUAL "RapidJSON") build_rapidjson() + elseif("${DEPENDENCY_NAME}" STREQUAL "DataSketches") + build_datasketches() elseif("${DEPENDENCY_NAME}" STREQUAL "zstd") build_zstd() elseif("${DEPENDENCY_NAME}" STREQUAL "Snappy") @@ -883,6 +898,26 @@ macro(build_rapidjson) add_dependencies(RapidJSON rapidjson_ep) endmacro() +macro(build_datasketches) + message(STATUS "Building DataSketches from source") + set(DATASKETCHES_PREFIX "${CMAKE_CURRENT_BINARY_DIR}/datasketches_ep-install") + set(DATASKETCHES_INCLUDE_DIR "${DATASKETCHES_PREFIX}/include") + set(DATASKETCHES_CMAKE_ARGS ${EP_COMMON_CMAKE_ARGS} -DBUILD_TESTS=OFF + "-DCMAKE_INSTALL_PREFIX=${DATASKETCHES_PREFIX}") + + externalproject_add(datasketches_ep + ${EP_COMMON_OPTIONS} + URL ${DATASKETCHES_SOURCE_URL} + URL_HASH "SHA256=${PAIMON_DATASKETCHES_BUILD_SHA256_CHECKSUM}" + CMAKE_ARGS ${DATASKETCHES_CMAKE_ARGS}) + + file(MAKE_DIRECTORY "${DATASKETCHES_INCLUDE_DIR}") + add_library(DataSketches INTERFACE IMPORTED) + target_include_directories(DataSketches SYSTEM + INTERFACE "${DATASKETCHES_INCLUDE_DIR}") + add_dependencies(DataSketches datasketches_ep) +endmacro() + macro(build_fmt) message(STATUS "Building fmt from source") set(FMT_PREFIX "${CMAKE_CURRENT_BINARY_DIR}/fmt_ep-install") @@ -1862,6 +1897,7 @@ endmacro() resolve_dependency(fmt) resolve_dependency(RapidJSON) +resolve_dependency(DataSketches) paimon_apply_dependency_source_defaults() resolve_dependency(RE2) resolve_dependency(Snappy) diff --git a/include/paimon/defs.h b/include/paimon/defs.h index cc128d478..472b3cd5a 100644 --- a/include/paimon/defs.h +++ b/include/paimon/defs.h @@ -80,6 +80,14 @@ struct PAIMON_EXPORT Options { static const char DEFAULT_AGG_FUNCTION[]; /// IGNORE_RETRACT is "ignore-retract" static const char IGNORE_RETRACT[]; + /// NESTED_KEY is "nested-key" + static const char NESTED_KEY[]; + /// NESTED_KEY_NULL_STRATEGY is "nested-key-null-strategy" + static const char NESTED_KEY_NULL_STRATEGY[]; + /// NESTED_SEQUENCE_FIELD is "nested-sequence-field" + static const char NESTED_SEQUENCE_FIELD[]; + /// COUNT_LIMIT is "count-limit" + static const char COUNT_LIMIT[]; /// "distinct" - Distinct option for aggregate functions like listagg. Default value is false. /// Example: fields.f.distinct=true to deduplicate values during aggregation. static const char DISTINCT[]; diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index e94782a5c..cdb4fa890 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt @@ -32,6 +32,8 @@ set(PAIMON_COMMON_SRCS common/data/columnar/columnar_row.cpp common/data/columnar/columnar_row_ref.cpp common/data/decimal.cpp + common/data/generic_array.cpp + common/data/generic_map.cpp common/data/internal_row.cpp common/data/record_batch.cpp common/data/serializer/binary_row_serializer.cpp @@ -286,6 +288,12 @@ set(PAIMON_CORE_SRCS core/mergetree/compact/universal_compaction.cpp core/mergetree/compact/early_full_compaction.cpp core/mergetree/compact/aggregate/aggregate_merge_function.cpp + core/mergetree/compact/aggregate/field_aggregate_utils.cpp + core/mergetree/compact/aggregate/field_collect_agg.cpp + core/mergetree/compact/aggregate/field_merge_map_agg.cpp + core/mergetree/compact/aggregate/field_nested_update_agg.cpp + core/mergetree/compact/aggregate/field_roaring_bitmap64_agg.cpp + core/mergetree/compact/aggregate/field_sketch_agg.cpp core/mergetree/compact/aggregate/field_sum_agg.cpp core/mergetree/compact/interval_partition.cpp core/mergetree/compact/loser_tree.cpp @@ -416,6 +424,7 @@ add_paimon_lib(paimon xxhash Threads::Threads RapidJSON + DataSketches ${PAIMON_OBJECT_STORE_LINK_LIBS} STATIC_LINK_LIBS arrow @@ -426,6 +435,7 @@ add_paimon_lib(paimon xxhash Threads::Threads RapidJSON + DataSketches ${PAIMON_OBJECT_STORE_LINK_LIBS} SHARED_LINK_FLAGS ${PAIMON_VERSION_SCRIPT_FLAGS}) @@ -707,14 +717,19 @@ if(PAIMON_BUILD_TESTS) core/mergetree/compact/aggregate/aggregate_merge_function_test.cpp core/mergetree/compact/aggregate/field_aggregator_factory_test.cpp core/mergetree/compact/aggregate/field_bool_agg_test.cpp + core/mergetree/compact/aggregate/field_collect_agg_test.cpp core/mergetree/compact/aggregate/field_first_non_null_value_agg_test.cpp core/mergetree/compact/aggregate/field_first_value_agg_test.cpp core/mergetree/compact/aggregate/field_ignore_retract_agg_test.cpp core/mergetree/compact/aggregate/field_last_non_null_value_agg_test.cpp core/mergetree/compact/aggregate/field_last_value_agg_test.cpp core/mergetree/compact/aggregate/field_listagg_agg_test.cpp + core/mergetree/compact/aggregate/field_merge_map_agg_test.cpp core/mergetree/compact/aggregate/field_min_max_agg_test.cpp + core/mergetree/compact/aggregate/field_nested_update_agg_test.cpp core/mergetree/compact/aggregate/field_primary_key_agg_test.cpp + core/mergetree/compact/aggregate/field_roaring_bitmap64_agg_test.cpp + core/mergetree/compact/aggregate/field_sketch_agg_test.cpp core/mergetree/compact/aggregate/field_sum_agg_test.cpp core/mergetree/compact/deduplicate_merge_function_test.cpp core/mergetree/compact/first_row_merge_function_test.cpp @@ -827,6 +842,7 @@ if(PAIMON_BUILD_TESTS) STATIC_LINK_LIBS paimon_shared test_utils_static + DataSketches ${TEST_STATIC_LINK_LIBS} ${GTEST_LINK_TOOLCHAIN}) diff --git a/src/paimon/common/data/generic_array.cpp b/src/paimon/common/data/generic_array.cpp new file mode 100644 index 000000000..386f31376 --- /dev/null +++ b/src/paimon/common/data/generic_array.cpp @@ -0,0 +1,164 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * 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 + * + * 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 "paimon/common/data/generic_array.h" + +#include +#include + +#include "paimon/memory/bytes.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/status.h" + +namespace paimon { + +int32_t GenericArray::Size() const { + return static_cast(values_.size()); +} + +const VariantType& GenericArray::ValueAt(int32_t pos) const { + assert(pos >= 0 && static_cast(pos) < values_.size()); + return values_[pos]; +} + +bool GenericArray::IsNullAt(int32_t pos) const { + return DataDefine::IsVariantNull(ValueAt(pos)); +} + +bool GenericArray::GetBoolean(int32_t pos) const { + return DataDefine::GetVariantValue(ValueAt(pos)); +} + +char GenericArray::GetByte(int32_t pos) const { + return DataDefine::GetVariantValue(ValueAt(pos)); +} + +int16_t GenericArray::GetShort(int32_t pos) const { + return DataDefine::GetVariantValue(ValueAt(pos)); +} + +int32_t GenericArray::GetInt(int32_t pos) const { + return DataDefine::GetVariantValue(ValueAt(pos)); +} + +int32_t GenericArray::GetDate(int32_t pos) const { + return GetInt(pos); +} + +int64_t GenericArray::GetLong(int32_t pos) const { + return DataDefine::GetVariantValue(ValueAt(pos)); +} + +float GenericArray::GetFloat(int32_t pos) const { + return DataDefine::GetVariantValue(ValueAt(pos)); +} + +double GenericArray::GetDouble(int32_t pos) const { + return DataDefine::GetVariantValue(ValueAt(pos)); +} + +BinaryString GenericArray::GetString(int32_t pos) const { + const VariantType& value = ValueAt(pos); + if (const auto* string = DataDefine::GetVariantPtr(value)) { + return *string; + } + return BinaryString::FromString(std::string(DataDefine::GetStringView(value)), + GetDefaultPool().get()); +} + +std::string_view GenericArray::GetStringView(int32_t pos) const { + return DataDefine::GetStringView(ValueAt(pos)); +} + +Decimal GenericArray::GetDecimal(int32_t pos, int32_t precision, int32_t scale) const { + return DataDefine::GetVariantValue(ValueAt(pos)); +} + +Timestamp GenericArray::GetTimestamp(int32_t pos, int32_t precision) const { + return DataDefine::GetVariantValue(ValueAt(pos)); +} + +std::shared_ptr GenericArray::GetBinary(int32_t pos) const { + const VariantType& value = ValueAt(pos); + if (const auto* bytes = DataDefine::GetVariantPtr>(value)) { + return *bytes; + } + return std::shared_ptr(Bytes::AllocateBytes( + std::string(DataDefine::GetStringView(value)), GetDefaultPool().get())); +} + +std::shared_ptr GenericArray::GetArray(int32_t pos) const { + return DataDefine::GetVariantValue>(ValueAt(pos)); +} + +std::shared_ptr GenericArray::GetMap(int32_t pos) const { + return DataDefine::GetVariantValue>(ValueAt(pos)); +} + +std::shared_ptr GenericArray::GetRow(int32_t pos, int32_t num_fields) const { + return DataDefine::GetVariantValue>(ValueAt(pos)); +} + +template +Result> GenericArray::ToPrimitiveArray() const { + std::vector result; + result.reserve(values_.size()); + for (const VariantType& value : values_) { + if (DataDefine::IsVariantNull(value)) { + return Status::Invalid("Primitive array must not contain a null value."); + } + result.push_back(DataDefine::GetVariantValue(value)); + } + return result; +} + +Result> GenericArray::ToBooleanArray() const { + std::vector result; + result.reserve(values_.size()); + for (const VariantType& value : values_) { + if (DataDefine::IsVariantNull(value)) { + return Status::Invalid("Primitive array must not contain a null value."); + } + result.push_back(DataDefine::GetVariantValue(value)); + } + return result; +} + +Result> GenericArray::ToByteArray() const { + return ToPrimitiveArray(); +} + +Result> GenericArray::ToShortArray() const { + return ToPrimitiveArray(); +} + +Result> GenericArray::ToIntArray() const { + return ToPrimitiveArray(); +} + +Result> GenericArray::ToLongArray() const { + return ToPrimitiveArray(); +} + +Result> GenericArray::ToFloatArray() const { + return ToPrimitiveArray(); +} + +Result> GenericArray::ToDoubleArray() const { + return ToPrimitiveArray(); +} + +} // namespace paimon diff --git a/src/paimon/common/data/generic_array.h b/src/paimon/common/data/generic_array.h new file mode 100644 index 000000000..21aee93ab --- /dev/null +++ b/src/paimon/common/data/generic_array.h @@ -0,0 +1,79 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * 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 + * + * 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 +#include +#include + +#include "paimon/common/data/data_define.h" +#include "paimon/common/data/internal_array.h" + +namespace paimon { + +/// GenericArray is a generic implementation of `InternalArray` backed by a vector of VariantType. + +/// @note Holders keep source arrays alive for non-owning string and binary elements. That is not +/// enough when a source is itself a view, in which case the caller must keep the data alive. +class GenericArray : public InternalArray { + public: + /// Create an array from materialized values and optional source-array holders. + /// + /// @param values Values stored in the array. + /// @param holders Source arrays retained for non-owning string and binary values. + explicit GenericArray(std::vector values, + std::vector> holders = {}) + : values_(std::move(values)), holders_(std::move(holders)) {} + + int32_t Size() const override; + bool IsNullAt(int32_t pos) const override; + bool GetBoolean(int32_t pos) const override; + char GetByte(int32_t pos) const override; + int16_t GetShort(int32_t pos) const override; + int32_t GetInt(int32_t pos) const override; + int32_t GetDate(int32_t pos) const override; + int64_t GetLong(int32_t pos) const override; + float GetFloat(int32_t pos) const override; + double GetDouble(int32_t pos) const override; + BinaryString GetString(int32_t pos) const override; + std::string_view GetStringView(int32_t pos) const override; + Decimal GetDecimal(int32_t pos, int32_t precision, int32_t scale) const override; + Timestamp GetTimestamp(int32_t pos, int32_t precision) const override; + std::shared_ptr GetBinary(int32_t pos) const override; + std::shared_ptr GetArray(int32_t pos) const override; + std::shared_ptr GetMap(int32_t pos) const override; + std::shared_ptr GetRow(int32_t pos, int32_t num_fields) const override; + + Result> ToBooleanArray() const override; + Result> ToByteArray() const override; + Result> ToShortArray() const override; + Result> ToIntArray() const override; + Result> ToLongArray() const override; + Result> ToFloatArray() const override; + Result> ToDoubleArray() const override; + + private: + template + Result> ToPrimitiveArray() const; + + const VariantType& ValueAt(int32_t pos) const; + + std::vector values_; + std::vector> holders_; +}; + +} // namespace paimon diff --git a/src/paimon/common/data/generic_map.cpp b/src/paimon/common/data/generic_map.cpp new file mode 100644 index 000000000..60aa3700b --- /dev/null +++ b/src/paimon/common/data/generic_map.cpp @@ -0,0 +1,42 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * 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 + * + * 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 "paimon/common/data/generic_map.h" + +#include +#include + +namespace paimon { + +GenericMap::GenericMap(std::shared_ptr key_array, + std::shared_ptr value_array) + : key_array_(std::move(key_array)), value_array_(std::move(value_array)) { + assert(key_array_ && value_array_ && key_array_->Size() == value_array_->Size()); +} + +int32_t GenericMap::Size() const { + return key_array_->Size(); +} + +std::shared_ptr GenericMap::KeyArray() const { + return key_array_; +} + +std::shared_ptr GenericMap::ValueArray() const { + return value_array_; +} + +} // namespace paimon diff --git a/src/paimon/common/data/generic_map.h b/src/paimon/common/data/generic_map.h new file mode 100644 index 000000000..43ad75304 --- /dev/null +++ b/src/paimon/common/data/generic_map.h @@ -0,0 +1,44 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * 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 + * + * 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 + +#include "paimon/common/data/internal_map.h" + +namespace paimon { + +/// An InternalMap backed by separate key and value arrays. +class GenericMap : public InternalMap { + public: + /// Create a map from equally sized key and value arrays. + /// + /// @param key_array Keys stored in the map. + /// @param value_array Values stored in the map. + GenericMap(std::shared_ptr key_array, + std::shared_ptr value_array); + + int32_t Size() const override; + std::shared_ptr KeyArray() const override; + std::shared_ptr ValueArray() const override; + + private: + std::shared_ptr key_array_; + std::shared_ptr value_array_; +}; + +} // namespace paimon diff --git a/src/paimon/common/defs.cpp b/src/paimon/common/defs.cpp index 0df896a99..01b99c474 100644 --- a/src/paimon/common/defs.cpp +++ b/src/paimon/common/defs.cpp @@ -23,6 +23,10 @@ const char Options::FIELDS_PREFIX[] = "fields"; const char Options::AGG_FUNCTION[] = "aggregate-function"; const char Options::DEFAULT_AGG_FUNCTION[] = "default-aggregate-function"; const char Options::IGNORE_RETRACT[] = "ignore-retract"; +const char Options::NESTED_KEY[] = "nested-key"; +const char Options::NESTED_KEY_NULL_STRATEGY[] = "nested-key-null-strategy"; +const char Options::NESTED_SEQUENCE_FIELD[] = "nested-sequence-field"; +const char Options::COUNT_LIMIT[] = "count-limit"; const char Options::DISTINCT[] = "distinct"; const char Options::LIST_AGG_DELIMITER[] = "list-agg-delimiter"; const char Options::SEQUENCE_GROUP[] = "sequence-group"; diff --git a/src/paimon/core/core_options.cpp b/src/paimon/core/core_options.cpp index 252590607..fff407154 100644 --- a/src/paimon/core/core_options.cpp +++ b/src/paimon/core/core_options.cpp @@ -1307,6 +1307,61 @@ Result CoreOptions::FieldAggIgnoreRetract(const std::string& field_name) c return field_agg_ignore_retract; } +Result> CoreOptions::FieldNestedUpdateAggNestedKey( + const std::string& field_name) const { + ConfigParser parser(impl_->raw_options); + std::vector nested_key; + std::string key = std::string(Options::FIELDS_PREFIX) + "." + field_name + "." + + std::string(Options::NESTED_KEY); + PAIMON_RETURN_NOT_OK( + parser.ParseList(key, Options::FIELDS_SEPARATOR, &nested_key, true)); + return nested_key; +} + +Result CoreOptions::FieldNestedUpdateAggNestedKeyNullStrategy( + const std::string& field_name) const { + ConfigParser parser(impl_->raw_options); + std::string strategy = "merge"; + std::string key = std::string(Options::FIELDS_PREFIX) + "." + field_name + "." + + std::string(Options::NESTED_KEY_NULL_STRATEGY); + PAIMON_RETURN_NOT_OK(parser.Parse(key, &strategy)); + if (strategy == "merge") { + return NestedKeyNullStrategy::MERGE; + } + if (strategy == "ignore") { + return NestedKeyNullStrategy::IGNORE; + } + if (strategy == "error") { + return NestedKeyNullStrategy::ERROR; + } + return Status::Invalid(fmt::format( + "Invalid Config [{}: {}], supported values are merge, ignore and error", key, strategy)); +} + +Result> CoreOptions::FieldNestedUpdateAggNestedSequenceField( + const std::string& field_name) const { + ConfigParser parser(impl_->raw_options); + std::vector sequence_fields; + std::string key = std::string(Options::FIELDS_PREFIX) + "." + field_name + "." + + std::string(Options::NESTED_SEQUENCE_FIELD); + PAIMON_RETURN_NOT_OK( + parser.ParseList(key, Options::FIELDS_SEPARATOR, &sequence_fields, true)); + return sequence_fields; +} + +Result CoreOptions::FieldNestedUpdateAggCountLimit(const std::string& field_name) const { + ConfigParser parser(impl_->raw_options); + int32_t count_limit = std::numeric_limits::max(); + std::string key = std::string(Options::FIELDS_PREFIX) + "." + field_name + "." + + std::string(Options::COUNT_LIMIT); + PAIMON_RETURN_NOT_OK(parser.Parse(key, &count_limit)); + if (count_limit < 0) { + return Status::Invalid( + fmt::format("Invalid Config [{}: {}], must not be negative", key, count_limit)); + } + return count_limit; +} + Result CoreOptions::FieldListAggDelimiter(const std::string& field_name) const { ConfigParser parser(impl_->raw_options); std::string delimiter = ","; diff --git a/src/paimon/core/core_options.h b/src/paimon/core/core_options.h index 860ae67f4..a301c1521 100644 --- a/src/paimon/core/core_options.h +++ b/src/paimon/core/core_options.h @@ -57,6 +57,13 @@ class PAIMON_EXPORT CoreOptions { SNAPSHOT, }; + /// Defines how nested_update handles null values in nested keys. + enum class NestedKeyNullStrategy { + MERGE, + IGNORE, + ERROR, + }; + static Result FromMap( const std::map& options_map, const std::shared_ptr& specified_file_system = nullptr, @@ -137,6 +144,29 @@ class PAIMON_EXPORT CoreOptions { std::optional GetFieldsDefaultFunc() const; Result> GetFieldAggFunc(const std::string& field_name) const; Result FieldAggIgnoreRetract(const std::string& field_name) const; + /// Return nested key fields configured for a nested_update field. + /// + /// @param field_name Name of the table field. + /// @return Configured nested key field names, or an error Status. + Result> FieldNestedUpdateAggNestedKey( + const std::string& field_name) const; + /// Return the null-key strategy configured for a nested_update field. + /// + /// @param field_name Name of the table field. + /// @return The configured null-key strategy, or an error Status. + Result FieldNestedUpdateAggNestedKeyNullStrategy( + const std::string& field_name) const; + /// Return sequence fields configured for a nested_update field. + /// + /// @param field_name Name of the table field. + /// @return Configured nested sequence field names, or an error Status. + Result> FieldNestedUpdateAggNestedSequenceField( + const std::string& field_name) const; + /// Return the maximum number of rows retained by a nested_update field. + /// + /// @param field_name Name of the table field. + /// @return The configured row count limit, or an error Status. + Result FieldNestedUpdateAggCountLimit(const std::string& field_name) const; Result FieldListAggDelimiter(const std::string& field_name) const; Result FieldCollectAggDistinct(const std::string& field_name) const; diff --git a/src/paimon/core/mergetree/compact/aggregate/aggregate_merge_function.cpp b/src/paimon/core/mergetree/compact/aggregate/aggregate_merge_function.cpp index 2230c80d1..9ec9f0528 100644 --- a/src/paimon/core/mergetree/compact/aggregate/aggregate_merge_function.cpp +++ b/src/paimon/core/mergetree/compact/aggregate/aggregate_merge_function.cpp @@ -85,7 +85,8 @@ Status AggregateMergeFunction::Add(KeyValue&& kv) { PAIMON_ASSIGN_OR_RAISE(merged_field, aggregators_[i]->Retract(accumulator, input_field)); } else { - merged_field = aggregators_[i]->Agg(accumulator, input_field); + PAIMON_ASSIGN_OR_RAISE(merged_field, + aggregators_[i]->AggResult(accumulator, input_field)); } row_->SetField(i, merged_field); } diff --git a/src/paimon/core/mergetree/compact/aggregate/field_aggregate_utils.cpp b/src/paimon/core/mergetree/compact/aggregate/field_aggregate_utils.cpp new file mode 100644 index 000000000..8ffca841d --- /dev/null +++ b/src/paimon/core/mergetree/compact/aggregate/field_aggregate_utils.cpp @@ -0,0 +1,237 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * 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 + * + * 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 "paimon/core/mergetree/compact/aggregate/field_aggregate_utils.h" + +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/util/checked_cast.h" +#include "fmt/format.h" +#include "paimon/common/data/data_getters.h" +#include "paimon/common/data/internal_array.h" +#include "paimon/common/data/internal_map.h" +#include "paimon/common/data/internal_row.h" +#include "paimon/common/utils/date_time_utils.h" +#include "paimon/common/utils/fields_comparator.h" +#include "paimon/memory/bytes.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/status.h" + +namespace paimon { +namespace { + +Result EqualGetters(const DataGetters& lhs, int32_t lhs_pos, const DataGetters& rhs, + int32_t rhs_pos, const std::shared_ptr& type) { + PAIMON_ASSIGN_OR_RAISE(VariantType lhs_value, + FieldAggregateUtils::GetValue(lhs, lhs_pos, type)); + PAIMON_ASSIGN_OR_RAISE(VariantType rhs_value, + FieldAggregateUtils::GetValue(rhs, rhs_pos, type)); + return FieldAggregateUtils::Equals(lhs_value, rhs_value, type); +} + +Result EqualRows(const std::shared_ptr& lhs, + const std::shared_ptr& rhs, + const std::shared_ptr& type) { + if (!lhs || !rhs || lhs->GetFieldCount() != rhs->GetFieldCount() || + lhs->GetFieldCount() != type->num_fields()) { + return lhs == rhs; + } + for (int32_t i = 0; i < type->num_fields(); ++i) { + PAIMON_ASSIGN_OR_RAISE(bool equal, EqualGetters(*lhs, i, *rhs, i, type->field(i)->type())); + if (!equal) { + return false; + } + } + return true; +} + +Result EqualArrays(const std::shared_ptr& lhs, + const std::shared_ptr& rhs, + const std::shared_ptr& type) { + if (!lhs || !rhs || lhs->Size() != rhs->Size()) { + return lhs == rhs; + } + for (int32_t i = 0; i < lhs->Size(); ++i) { + PAIMON_ASSIGN_OR_RAISE(bool equal, EqualGetters(*lhs, i, *rhs, i, type->value_type())); + if (!equal) { + return false; + } + } + return true; +} + +Result EqualMaps(const std::shared_ptr& lhs, + const std::shared_ptr& rhs, + const std::shared_ptr& type) { + if (!lhs || !rhs || lhs->Size() != rhs->Size()) { + return lhs == rhs; + } + std::shared_ptr lhs_keys = lhs->KeyArray(); + std::shared_ptr lhs_values = lhs->ValueArray(); + std::shared_ptr rhs_keys = rhs->KeyArray(); + std::shared_ptr rhs_values = rhs->ValueArray(); + std::vector matched(rhs->Size(), false); + for (int32_t i = 0; i < lhs->Size(); ++i) { + bool found = false; + for (int32_t j = 0; j < rhs->Size(); ++j) { + if (matched[j]) { + continue; + } + PAIMON_ASSIGN_OR_RAISE(bool key_equal, + EqualGetters(*lhs_keys, i, *rhs_keys, j, type->key_type())); + if (!key_equal) { + continue; + } + PAIMON_ASSIGN_OR_RAISE(bool value_equal, + EqualGetters(*lhs_values, i, *rhs_values, j, type->item_type())); + if (value_equal) { + matched[j] = true; + found = true; + break; + } + } + if (!found) { + return false; + } + } + return true; +} + +} // namespace + +VariantType FieldAggregateUtils::OwnedBinary(const VariantType& value) { + if (DataDefine::GetVariantPtr>(value)) { + return value; + } + std::string_view view = DataDefine::GetStringView(value); + pooled_unique_ptr owned = Bytes::AllocateBytes(view.size(), GetDefaultPool().get()); + if (!view.empty()) { + std::memcpy(owned->data(), view.data(), view.size()); + } + return VariantType(std::shared_ptr(std::move(owned))); +} + +Result FieldAggregateUtils::GetValue(const DataGetters& getters, int32_t pos, + const std::shared_ptr& type) { + if (getters.IsNullAt(pos)) { + return VariantType(NullType()); + } + switch (type->id()) { + case arrow::Type::BOOL: + return VariantType(getters.GetBoolean(pos)); + case arrow::Type::INT8: + return VariantType(getters.GetByte(pos)); + case arrow::Type::INT16: + return VariantType(getters.GetShort(pos)); + case arrow::Type::DATE32: + return VariantType(getters.GetDate(pos)); + case arrow::Type::INT32: + return VariantType(getters.GetInt(pos)); + case arrow::Type::INT64: + return VariantType(getters.GetLong(pos)); + case arrow::Type::FLOAT: + return VariantType(getters.GetFloat(pos)); + case arrow::Type::DOUBLE: + return VariantType(getters.GetDouble(pos)); + case arrow::Type::STRING: + case arrow::Type::BINARY: + return VariantType(getters.GetStringView(pos)); + case arrow::Type::TIMESTAMP: { + std::shared_ptr timestamp_type = + arrow::internal::checked_pointer_cast(type); + return VariantType( + getters.GetTimestamp(pos, DateTimeUtils::GetPrecisionFromType(timestamp_type))); + } + case arrow::Type::DECIMAL128: { + const auto* decimal_type = + arrow::internal::checked_cast(type.get()); + return VariantType( + getters.GetDecimal(pos, decimal_type->precision(), decimal_type->scale())); + } + case arrow::Type::LIST: + return VariantType(getters.GetArray(pos)); + case arrow::Type::MAP: + return VariantType(getters.GetMap(pos)); + case arrow::Type::STRUCT: + return VariantType(getters.GetRow(pos, type->num_fields())); + default: + return Status::Invalid( + fmt::format("type {} is not supported by field aggregation", type->ToString())); + } +} + +Result FieldAggregateUtils::Equals(const VariantType& lhs, const VariantType& rhs, + const std::shared_ptr& type) { + bool lhs_null = DataDefine::IsVariantNull(lhs); + bool rhs_null = DataDefine::IsVariantNull(rhs); + if (lhs_null || rhs_null) { + return lhs_null && rhs_null; + } + switch (type->id()) { + case arrow::Type::BOOL: + return DataDefine::GetVariantValue(lhs) == DataDefine::GetVariantValue(rhs); + case arrow::Type::INT8: + return DataDefine::GetVariantValue(lhs) == DataDefine::GetVariantValue(rhs); + case arrow::Type::INT16: + return DataDefine::GetVariantValue(lhs) == + DataDefine::GetVariantValue(rhs); + case arrow::Type::DATE32: + case arrow::Type::INT32: + return DataDefine::GetVariantValue(lhs) == + DataDefine::GetVariantValue(rhs); + case arrow::Type::INT64: + return DataDefine::GetVariantValue(lhs) == + DataDefine::GetVariantValue(rhs); + case arrow::Type::FLOAT: + return FieldsComparator::CompareFloatingPoint( + DataDefine::GetVariantValue(lhs), + DataDefine::GetVariantValue(rhs)) == 0; + case arrow::Type::DOUBLE: + return FieldsComparator::CompareFloatingPoint( + DataDefine::GetVariantValue(lhs), + DataDefine::GetVariantValue(rhs)) == 0; + case arrow::Type::STRING: + case arrow::Type::BINARY: + return DataDefine::GetStringView(lhs) == DataDefine::GetStringView(rhs); + case arrow::Type::TIMESTAMP: + return DataDefine::GetVariantValue(lhs) == + DataDefine::GetVariantValue(rhs); + case arrow::Type::DECIMAL128: + return DataDefine::GetVariantValue(lhs) == + DataDefine::GetVariantValue(rhs); + case arrow::Type::STRUCT: + return EqualRows(DataDefine::GetVariantValue>(lhs), + DataDefine::GetVariantValue>(rhs), + arrow::internal::checked_pointer_cast(type)); + case arrow::Type::LIST: + return EqualArrays(DataDefine::GetVariantValue>(lhs), + DataDefine::GetVariantValue>(rhs), + arrow::internal::checked_pointer_cast(type)); + case arrow::Type::MAP: + return EqualMaps(DataDefine::GetVariantValue>(lhs), + DataDefine::GetVariantValue>(rhs), + arrow::internal::checked_pointer_cast(type)); + default: + return Status::Invalid( + fmt::format("type {} is not supported by field aggregation", type->ToString())); + } +} + +} // namespace paimon diff --git a/src/paimon/core/mergetree/compact/aggregate/field_aggregate_utils.h b/src/paimon/core/mergetree/compact/aggregate/field_aggregate_utils.h new file mode 100644 index 000000000..dcf404c04 --- /dev/null +++ b/src/paimon/core/mergetree/compact/aggregate/field_aggregate_utils.h @@ -0,0 +1,62 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * 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 + * + * 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 + +#include "arrow/type_fwd.h" +#include "paimon/common/data/data_define.h" +#include "paimon/result.h" + +namespace paimon { + +class Bytes; +class DataGetters; + +/// Helpers for extracting and comparing field aggregation values. +class FieldAggregateUtils { + public: + FieldAggregateUtils() = delete; + ~FieldAggregateUtils() = delete; + + /// Return a binary value which owns its buffer. A merged row hands out binary fields as views + /// into buffers it releases once the field is overwritten. + /// + /// @param value Binary value, either owning or a view. + /// @return The value itself when it already owns its buffer, otherwise an owning copy. + static VariantType OwnedBinary(const VariantType& value); + + /// Extract a value from a typed field. + /// + /// @param getters Source containing the field. + /// @param pos Field position in the source. + /// @param type Logical type of the field. + /// @return The extracted value, or an error Status. + static Result GetValue(const DataGetters& getters, int32_t pos, + const std::shared_ptr& type); + + /// Compare two values using their logical type. + /// + /// @param lhs Left value. + /// @param rhs Right value. + /// @param type Logical type of both values. + /// @return Whether the values are equal, or an error Status. + static Result Equals(const VariantType& lhs, const VariantType& rhs, + const std::shared_ptr& type); +}; + +} // namespace paimon diff --git a/src/paimon/core/mergetree/compact/aggregate/field_aggregator.h b/src/paimon/core/mergetree/compact/aggregate/field_aggregator.h index 4100a008d..dbecf287d 100644 --- a/src/paimon/core/mergetree/compact/aggregate/field_aggregator.h +++ b/src/paimon/core/mergetree/compact/aggregate/field_aggregator.h @@ -35,6 +35,17 @@ class FieldAggregator { virtual VariantType Agg(const VariantType& accumulator, const VariantType& input_field) = 0; + /// Fallible aggregation entry point used by merge functions. Aggregators which deserialize + /// external representations override this method to report invalid input without throwing. + /// + /// @param accumulator Current aggregated value. + /// @param input_field Value to merge into the accumulator. + /// @return The merged value, or an error Status. + virtual Result AggResult(const VariantType& accumulator, + const VariantType& input_field) { + return Agg(accumulator, input_field); + } + /// reset the aggregator to a clean start state. virtual void Reset() {} @@ -43,6 +54,16 @@ class FieldAggregator { return Agg(input_field, accumulator); } + /// Fallible reversed aggregation entry point used by merge functions. + /// + /// @param accumulator Current aggregated value. + /// @param input_field Older value to merge into the accumulator. + /// @return The merged value, or an error Status. + virtual Result AggReversedResult(const VariantType& accumulator, + const VariantType& input_field) { + return AggResult(input_field, accumulator); + } + virtual Result Retract(const VariantType& accumulator, const VariantType& input_field) const { return Status::Invalid(fmt::format( diff --git a/src/paimon/core/mergetree/compact/aggregate/field_aggregator_factory.h b/src/paimon/core/mergetree/compact/aggregate/field_aggregator_factory.h index 30f25a94f..d59002439 100644 --- a/src/paimon/core/mergetree/compact/aggregate/field_aggregator_factory.h +++ b/src/paimon/core/mergetree/compact/aggregate/field_aggregator_factory.h @@ -24,6 +24,7 @@ #include "paimon/core/mergetree/compact/aggregate/field_aggregator.h" #include "paimon/core/mergetree/compact/aggregate/field_bool_and_agg.h" #include "paimon/core/mergetree/compact/aggregate/field_bool_or_agg.h" +#include "paimon/core/mergetree/compact/aggregate/field_collect_agg.h" #include "paimon/core/mergetree/compact/aggregate/field_first_non_null_value_agg.h" #include "paimon/core/mergetree/compact/aggregate/field_first_value_agg.h" #include "paimon/core/mergetree/compact/aggregate/field_ignore_retract_agg.h" @@ -31,8 +32,12 @@ #include "paimon/core/mergetree/compact/aggregate/field_last_value_agg.h" #include "paimon/core/mergetree/compact/aggregate/field_listagg_agg.h" #include "paimon/core/mergetree/compact/aggregate/field_max_agg.h" +#include "paimon/core/mergetree/compact/aggregate/field_merge_map_agg.h" #include "paimon/core/mergetree/compact/aggregate/field_min_agg.h" +#include "paimon/core/mergetree/compact/aggregate/field_nested_update_agg.h" #include "paimon/core/mergetree/compact/aggregate/field_primary_key_agg.h" +#include "paimon/core/mergetree/compact/aggregate/field_roaring_bitmap64_agg.h" +#include "paimon/core/mergetree/compact/aggregate/field_sketch_agg.h" #include "paimon/core/mergetree/compact/aggregate/field_sum_agg.h" #include "paimon/result.h" #include "paimon/status.h" @@ -75,6 +80,24 @@ class FieldAggregatorFactory { } else if (str_agg == FieldListaggAgg::NAME) { PAIMON_ASSIGN_OR_RAISE(field_aggregator, FieldListaggAgg::Create(field_type, options, field_name)); + } else if (str_agg == FieldCollectAgg::NAME) { + PAIMON_ASSIGN_OR_RAISE(field_aggregator, + FieldCollectAgg::Create(field_type, options, field_name)); + } else if (str_agg == FieldMergeMapAgg::NAME) { + PAIMON_ASSIGN_OR_RAISE(field_aggregator, + FieldMergeMapAgg::Create(field_type, field_name)); + } else if (str_agg == FieldNestedUpdateAgg::NAME) { + PAIMON_ASSIGN_OR_RAISE(field_aggregator, + FieldNestedUpdateAgg::Create(field_type, options, field_name)); + } else if (str_agg == FieldRoaringBitmap64Agg::NAME) { + PAIMON_ASSIGN_OR_RAISE(field_aggregator, + FieldRoaringBitmap64Agg::Create(field_type, field_name)); + } else if (str_agg == FieldHllSketchAgg::NAME) { + PAIMON_ASSIGN_OR_RAISE(field_aggregator, + FieldHllSketchAgg::Create(field_type, field_name)); + } else if (str_agg == FieldThetaSketchAgg::NAME) { + PAIMON_ASSIGN_OR_RAISE(field_aggregator, + FieldThetaSketchAgg::Create(field_type, field_name)); } else { return Status::Invalid(fmt::format( "Use unsupported aggregation {} or spell aggregate function incorrectly!", diff --git a/src/paimon/core/mergetree/compact/aggregate/field_aggregator_factory_test.cpp b/src/paimon/core/mergetree/compact/aggregate/field_aggregator_factory_test.cpp index c570dc874..38e655cc3 100644 --- a/src/paimon/core/mergetree/compact/aggregate/field_aggregator_factory_test.cpp +++ b/src/paimon/core/mergetree/compact/aggregate/field_aggregator_factory_test.cpp @@ -17,7 +17,9 @@ #include "paimon/core/mergetree/compact/aggregate/field_aggregator_factory.h" #include +#include +#include "arrow/api.h" #include "arrow/type_fwd.h" #include "gtest/gtest.h" #include "paimon/testing/utils/testharness.h" @@ -124,4 +126,31 @@ TEST(FieldAggregatorFactoryTest, TestRemoveRecordOnDeleteConflictsWithIgnoreRetr "conflicting behavior"); } +TEST(FieldAggregatorFactoryTest, CreatesJavaCompatibleAggregators) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({})); + std::shared_ptr nested_type = + arrow::list(arrow::struct_({arrow::field("id", arrow::int32())})); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr collect, + FieldAggregatorFactory::CreateFieldAggregator( + "f", arrow::list(arrow::int32()), "collect", options)); + ASSERT_TRUE(dynamic_cast(collect.get())); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr merge_map, + FieldAggregatorFactory::CreateFieldAggregator( + "f", arrow::map(arrow::int32(), arrow::int32()), "merge_map", options)); + ASSERT_TRUE(dynamic_cast(merge_map.get())); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr nested_update, + FieldAggregatorFactory::CreateFieldAggregator("f", nested_type, "nested_update", options)); + ASSERT_TRUE(dynamic_cast(nested_update.get())); + + for (const char* name : {"rbm64", "hll_sketch", "theta_sketch"}) { + ASSERT_OK_AND_ASSIGN( + std::unique_ptr aggregator, + FieldAggregatorFactory::CreateFieldAggregator("f", arrow::binary(), name, options)); + ASSERT_EQ(name, aggregator->GetName()); + } +} + } // namespace paimon::test diff --git a/src/paimon/core/mergetree/compact/aggregate/field_collect_agg.cpp b/src/paimon/core/mergetree/compact/aggregate/field_collect_agg.cpp new file mode 100644 index 000000000..c50f1c16c --- /dev/null +++ b/src/paimon/core/mergetree/compact/aggregate/field_collect_agg.cpp @@ -0,0 +1,175 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * 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 + * + * 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 "paimon/core/mergetree/compact/aggregate/field_collect_agg.h" + +#include + +#include "arrow/api.h" +#include "arrow/util/checked_cast.h" +#include "fmt/format.h" +#include "paimon/common/data/generic_array.h" +#include "paimon/common/data/internal_array.h" +#include "paimon/core/core_options.h" +#include "paimon/core/mergetree/compact/aggregate/field_aggregate_utils.h" +#include "paimon/status.h" + +namespace paimon { +namespace { + +Result Contains(const std::vector& values, const VariantType& candidate, + const std::shared_ptr& element_type) { + for (const VariantType& value : values) { + PAIMON_ASSIGN_OR_RAISE(bool equal, + FieldAggregateUtils::Equals(value, candidate, element_type)); + if (equal) { + return true; + } + } + return false; +} + +Status AppendArray(const std::shared_ptr& array, + const std::shared_ptr& element_type, bool distinct, + std::vector* values) { + if (!array) { + return Status::OK(); + } + for (int32_t i = 0; i < array->Size(); ++i) { + PAIMON_ASSIGN_OR_RAISE(VariantType value, + FieldAggregateUtils::GetValue(*array, i, element_type)); + if (distinct) { + PAIMON_ASSIGN_OR_RAISE(bool contains, Contains(*values, value, element_type)); + if (contains) { + continue; + } + } + values->push_back(std::move(value)); + } + return Status::OK(); +} + +} // namespace + +Result> FieldCollectAgg::Create( + const std::shared_ptr& field_type, const CoreOptions& options, + const std::string& field_name) { + if (field_type->id() != arrow::Type::LIST) { + return Status::Invalid( + fmt::format("invalid field type {} for field '{}' of {}, supposed to be array", + field_type->ToString(), field_name, NAME)); + } + std::shared_ptr list_type = + arrow::internal::checked_pointer_cast(field_type); + PAIMON_ASSIGN_OR_RAISE(bool distinct, options.FieldCollectAggDistinct(field_name)); + return std::unique_ptr( + new FieldCollectAgg(field_type, list_type->value_type(), distinct)); +} + +VariantType FieldCollectAgg::Agg(const VariantType& accumulator, const VariantType& input_field) { + Result result = AggImpl(accumulator, input_field); + return result.ok() ? std::move(result).value() : accumulator; +} + +Result FieldCollectAgg::AggResult(const VariantType& accumulator, + const VariantType& input_field) { + return AggImpl(accumulator, input_field); +} + +VariantType FieldCollectAgg::AggReversed(const VariantType& accumulator, + const VariantType& input_field) { + return Agg(accumulator, input_field); +} + +Result FieldCollectAgg::AggReversedResult(const VariantType& accumulator, + const VariantType& input_field) { + return AggImpl(accumulator, input_field); +} + +Result FieldCollectAgg::AggImpl(const VariantType& accumulator, + const VariantType& input_field) const { + bool accumulator_null = DataDefine::IsVariantNull(accumulator); + bool input_null = DataDefine::IsVariantNull(input_field); + if (accumulator_null && input_null) { + return VariantType(NullType()); + } + if (!distinct_ && (accumulator_null || input_null)) { + return accumulator_null ? input_field : accumulator; + } + + std::shared_ptr accumulator_array = + accumulator_null ? nullptr + : DataDefine::GetVariantValue>(accumulator); + std::shared_ptr input_array = + input_null ? nullptr + : DataDefine::GetVariantValue>(input_field); + std::vector values; + if (accumulator_array) { + values.reserve(accumulator_array->Size() + (input_array ? input_array->Size() : 0)); + } + PAIMON_RETURN_NOT_OK(AppendArray(accumulator_array, element_type_, distinct_, &values)); + PAIMON_RETURN_NOT_OK(AppendArray(input_array, element_type_, distinct_, &values)); + std::vector> holders; + if (accumulator_array) { + holders.push_back(accumulator_array); + } + if (input_array) { + holders.push_back(input_array); + } + return VariantType(std::static_pointer_cast( + std::make_shared(std::move(values), std::move(holders)))); +} + +Result FieldCollectAgg::Retract(const VariantType& accumulator, + const VariantType& input_field) const { + if (DataDefine::IsVariantNull(accumulator) || DataDefine::IsVariantNull(input_field)) { + return accumulator; + } + auto accumulator_array = + DataDefine::GetVariantValue>(accumulator); + auto retract_array = DataDefine::GetVariantValue>(input_field); + if (retract_array->Size() == 0) { + return accumulator; + } + + std::vector retract_values; + PAIMON_RETURN_NOT_OK( + AppendArray(retract_array, element_type_, /*distinct=*/false, &retract_values)); + std::vector result_values; + result_values.reserve(accumulator_array->Size()); + for (int32_t i = 0; i < accumulator_array->Size(); ++i) { + PAIMON_ASSIGN_OR_RAISE(VariantType candidate, + FieldAggregateUtils::GetValue(*accumulator_array, i, element_type_)); + bool removed = false; + for (auto iter = retract_values.begin(); iter != retract_values.end(); ++iter) { + PAIMON_ASSIGN_OR_RAISE(bool equal, + FieldAggregateUtils::Equals(candidate, *iter, element_type_)); + if (equal) { + retract_values.erase(iter); + removed = true; + break; + } + } + if (!removed) { + result_values.push_back(std::move(candidate)); + } + } + return VariantType(std::static_pointer_cast(std::make_shared( + std::move(result_values), + std::vector>{accumulator_array, retract_array}))); +} + +} // namespace paimon diff --git a/src/paimon/core/mergetree/compact/aggregate/field_collect_agg.h b/src/paimon/core/mergetree/compact/aggregate/field_collect_agg.h new file mode 100644 index 000000000..430b24c7f --- /dev/null +++ b/src/paimon/core/mergetree/compact/aggregate/field_collect_agg.h @@ -0,0 +1,67 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * 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 + * + * 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 +#include + +#include "paimon/core/mergetree/compact/aggregate/field_aggregator.h" + +namespace paimon { + +class CoreOptions; + +/// Concatenates arrays, optionally removing duplicate elements. +class FieldCollectAgg : public FieldAggregator { + public: + static constexpr char NAME[] = "collect"; + + /// Create a collect aggregator for an array field. + /// + /// @param field_type Type of the aggregated field. + /// @param options Table options containing the distinct setting. + /// @param field_name Name of the aggregated field. + /// @return A collect aggregator, or an error Status. + static Result> Create( + const std::shared_ptr& field_type, const CoreOptions& options, + const std::string& field_name); + + VariantType Agg(const VariantType& accumulator, const VariantType& input_field) override; + Result AggResult(const VariantType& accumulator, + const VariantType& input_field) override; + VariantType AggReversed(const VariantType& accumulator, + const VariantType& input_field) override; + Result AggReversedResult(const VariantType& accumulator, + const VariantType& input_field) override; + Result Retract(const VariantType& accumulator, + const VariantType& input_field) const override; + + private: + FieldCollectAgg(const std::shared_ptr& field_type, + std::shared_ptr element_type, bool distinct) + : FieldAggregator(NAME, field_type), + element_type_(std::move(element_type)), + distinct_(distinct) {} + + Result AggImpl(const VariantType& accumulator, + const VariantType& input_field) const; + + std::shared_ptr element_type_; + bool distinct_; +}; + +} // namespace paimon diff --git a/src/paimon/core/mergetree/compact/aggregate/field_collect_agg_test.cpp b/src/paimon/core/mergetree/compact/aggregate/field_collect_agg_test.cpp new file mode 100644 index 000000000..5aef3eba7 --- /dev/null +++ b/src/paimon/core/mergetree/compact/aggregate/field_collect_agg_test.cpp @@ -0,0 +1,92 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * 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 + * + * 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 "paimon/core/mergetree/compact/aggregate/field_collect_agg.h" + +#include +#include + +#include "arrow/api.h" +#include "gtest/gtest.h" +#include "paimon/common/data/generic_array.h" +#include "paimon/common/data/serializer/binary_serializer_utils.h" +#include "paimon/core/core_options.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { +namespace { + +VariantType IntArray(std::vector values) { + return VariantType( + std::static_pointer_cast(std::make_shared(std::move(values)))); +} + +std::vector Values(const VariantType& value) { + auto array = DataDefine::GetVariantValue>(value); + std::vector values; + for (int32_t i = 0; i < array->Size(); ++i) { + values.push_back(array->GetInt(i)); + } + return values; +} + +Result> MakeCollectAgg(bool distinct) { + PAIMON_ASSIGN_OR_RAISE( + CoreOptions options, + CoreOptions::FromMap({{"fields.f.distinct", distinct ? "true" : "false"}})); + return FieldCollectAgg::Create(arrow::list(arrow::int32()), options, "f"); +} + +} // namespace + +TEST(FieldCollectAggTest, ConcatenatesWithoutReversing) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, MakeCollectAgg(false)); + VariantType left = IntArray({int32_t{1}, int32_t{2}}); + VariantType right = IntArray({int32_t{3}, int32_t{4}}); + + ASSERT_OK_AND_ASSIGN(VariantType result, agg->AggReversedResult(left, right)); + ASSERT_EQ((std::vector{1, 2, 3, 4}), Values(result)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr binary_result, + BinarySerializerUtils::WriteBinaryArray( + DataDefine::GetVariantValue>(result), + arrow::list(arrow::int32()), GetDefaultPool().get())); + ASSERT_EQ((std::vector{1, 2, 3, 4}), binary_result->ToIntArray().value()); + + ASSERT_OK_AND_ASSIGN(VariantType null_result, + agg->AggResult(VariantType(NullType()), VariantType(NullType()))); + ASSERT_TRUE(DataDefine::IsVariantNull(null_result)); +} + +TEST(FieldCollectAggTest, DistinctAndRetractOneOccurrence) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr distinct_agg, MakeCollectAgg(true)); + ASSERT_OK_AND_ASSIGN(VariantType distinct_result, + distinct_agg->AggResult(IntArray({int32_t{1}, int32_t{2}, int32_t{2}}), + IntArray({int32_t{2}, int32_t{3}}))); + ASSERT_EQ((std::vector{1, 2, 3}), Values(distinct_result)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, MakeCollectAgg(false)); + ASSERT_OK_AND_ASSIGN(VariantType retract_result, + agg->Retract(IntArray({int32_t{1}, int32_t{2}, int32_t{2}, int32_t{3}}), + IntArray({int32_t{2}}))); + ASSERT_EQ((std::vector{1, 2, 3}), Values(retract_result)); +} + +TEST(FieldCollectAggTest, RejectsNonArrayType) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({})); + ASSERT_NOK(FieldCollectAgg::Create(arrow::int32(), options, "f")); +} + +} // namespace paimon::test diff --git a/src/paimon/core/mergetree/compact/aggregate/field_ignore_retract_agg.h b/src/paimon/core/mergetree/compact/aggregate/field_ignore_retract_agg.h index fe41d82b4..06ad29592 100644 --- a/src/paimon/core/mergetree/compact/aggregate/field_ignore_retract_agg.h +++ b/src/paimon/core/mergetree/compact/aggregate/field_ignore_retract_agg.h @@ -34,6 +34,11 @@ class FieldIgnoreRetractAgg : public FieldAggregator { return agg_->Agg(accumulator, input_field); } + Result AggResult(const VariantType& accumulator, + const VariantType& input_field) override { + return agg_->AggResult(accumulator, input_field); + } + Result Retract(const VariantType& accumulator, const VariantType& input_field) const override { return accumulator; diff --git a/src/paimon/core/mergetree/compact/aggregate/field_ignore_retract_agg_test.cpp b/src/paimon/core/mergetree/compact/aggregate/field_ignore_retract_agg_test.cpp index dc81b77e7..3dd3c3a66 100644 --- a/src/paimon/core/mergetree/compact/aggregate/field_ignore_retract_agg_test.cpp +++ b/src/paimon/core/mergetree/compact/aggregate/field_ignore_retract_agg_test.cpp @@ -17,9 +17,16 @@ #include "paimon/core/mergetree/compact/aggregate/field_ignore_retract_agg.h" #include +#include +#include +#include +#include "arrow/api.h" #include "arrow/type_fwd.h" #include "gtest/gtest.h" +#include "paimon/common/data/generic_array.h" +#include "paimon/core/core_options.h" +#include "paimon/core/mergetree/compact/aggregate/field_collect_agg.h" #include "paimon/core/mergetree/compact/aggregate/field_sum_agg.h" #include "paimon/status.h" #include "paimon/testing/utils/testharness.h" @@ -64,4 +71,26 @@ TEST(FieldIgnoreRetractAggTest, TestNull) { } } +// matches Java, where the wrapper only overrides agg: reversed aggregation falls back to the base +// implementation and therefore bypasses the wrapped aggregator's own aggReversed override +TEST(FieldIgnoreRetractAggTest, ReversedAggBypassesWrappedOverride) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({})); + ASSERT_OK_AND_ASSIGN(std::unique_ptr collect_agg, + FieldCollectAgg::Create(arrow::list(arrow::int32()), options, "f0")); + auto agg = std::make_unique(std::move(collect_agg)); + + VariantType accumulator = VariantType(std::static_pointer_cast( + std::make_shared(std::vector{int32_t{1}, int32_t{2}}))); + VariantType input = VariantType(std::static_pointer_cast( + std::make_shared(std::vector{int32_t{3}, int32_t{4}}))); + + ASSERT_OK_AND_ASSIGN(VariantType result, agg->AggReversedResult(accumulator, input)); + auto values = DataDefine::GetVariantValue>(result); + ASSERT_EQ(4, values->Size()); + ASSERT_EQ(3, values->GetInt(0)); + ASSERT_EQ(4, values->GetInt(1)); + ASSERT_EQ(1, values->GetInt(2)); + ASSERT_EQ(2, values->GetInt(3)); +} + } // namespace paimon::test diff --git a/src/paimon/core/mergetree/compact/aggregate/field_merge_map_agg.cpp b/src/paimon/core/mergetree/compact/aggregate/field_merge_map_agg.cpp new file mode 100644 index 000000000..881b1b658 --- /dev/null +++ b/src/paimon/core/mergetree/compact/aggregate/field_merge_map_agg.cpp @@ -0,0 +1,157 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * 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 + * + * 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 "paimon/core/mergetree/compact/aggregate/field_merge_map_agg.h" + +#include +#include + +#include "arrow/api.h" +#include "arrow/util/checked_cast.h" +#include "fmt/format.h" +#include "paimon/common/data/generic_array.h" +#include "paimon/common/data/generic_map.h" +#include "paimon/core/mergetree/compact/aggregate/field_aggregate_utils.h" +#include "paimon/status.h" + +namespace paimon { +namespace { + +struct MapEntry { + VariantType key; + VariantType value; +}; + +Result FindKey(const std::vector& entries, const VariantType& key, + const std::shared_ptr& key_type) { + for (int32_t i = 0; i < static_cast(entries.size()); ++i) { + PAIMON_ASSIGN_OR_RAISE(bool equal, + FieldAggregateUtils::Equals(entries[i].key, key, key_type)); + if (equal) { + return i; + } + } + return -1; +} + +Status PutMap(const std::shared_ptr& map, + const std::shared_ptr& key_type, + const std::shared_ptr& value_type, std::vector* entries) { + std::shared_ptr keys = map->KeyArray(); + std::shared_ptr values = map->ValueArray(); + for (int32_t i = 0; i < map->Size(); ++i) { + PAIMON_ASSIGN_OR_RAISE(VariantType key, FieldAggregateUtils::GetValue(*keys, i, key_type)); + PAIMON_ASSIGN_OR_RAISE(VariantType value, + FieldAggregateUtils::GetValue(*values, i, value_type)); + PAIMON_ASSIGN_OR_RAISE(int32_t existing, FindKey(*entries, key, key_type)); + if (existing >= 0) { + (*entries)[existing].value = std::move(value); + } else { + entries->push_back(MapEntry{std::move(key), std::move(value)}); + } + } + return Status::OK(); +} + +VariantType MakeMap(std::vector entries, + std::vector> key_holders, + std::vector> value_holders) { + std::vector keys; + std::vector values; + keys.reserve(entries.size()); + values.reserve(entries.size()); + for (MapEntry& entry : entries) { + keys.push_back(std::move(entry.key)); + values.push_back(std::move(entry.value)); + } + std::shared_ptr key_array = + std::make_shared(std::move(keys), std::move(key_holders)); + std::shared_ptr value_array = + std::make_shared(std::move(values), std::move(value_holders)); + return std::static_pointer_cast( + std::make_shared(std::move(key_array), std::move(value_array))); +} + +} // namespace + +Result> FieldMergeMapAgg::Create( + const std::shared_ptr& field_type, const std::string& field_name) { + if (field_type->id() != arrow::Type::MAP) { + return Status::Invalid( + fmt::format("invalid field type {} for field '{}' of {}, supposed to be map", + field_type->ToString(), field_name, NAME)); + } + std::shared_ptr map_type = + arrow::internal::checked_pointer_cast(field_type); + return std::unique_ptr( + new FieldMergeMapAgg(field_type, map_type->key_type(), map_type->item_type())); +} + +VariantType FieldMergeMapAgg::Agg(const VariantType& accumulator, const VariantType& input_field) { + Result result = AggImpl(accumulator, input_field); + return result.ok() ? std::move(result).value() : accumulator; +} + +Result FieldMergeMapAgg::AggResult(const VariantType& accumulator, + const VariantType& input_field) { + return AggImpl(accumulator, input_field); +} + +Result FieldMergeMapAgg::AggImpl(const VariantType& accumulator, + const VariantType& input_field) const { + bool accumulator_null = DataDefine::IsVariantNull(accumulator); + bool input_null = DataDefine::IsVariantNull(input_field); + if (accumulator_null || input_null) { + return accumulator_null ? input_field : accumulator; + } + auto accumulator_map = DataDefine::GetVariantValue>(accumulator); + auto input_map = DataDefine::GetVariantValue>(input_field); + std::vector entries; + entries.reserve(accumulator_map->Size() + input_map->Size()); + PAIMON_RETURN_NOT_OK(PutMap(accumulator_map, key_type_, value_type_, &entries)); + PAIMON_RETURN_NOT_OK(PutMap(input_map, key_type_, value_type_, &entries)); + return MakeMap(std::move(entries), {accumulator_map->KeyArray(), input_map->KeyArray()}, + {accumulator_map->ValueArray(), input_map->ValueArray()}); +} + +Result FieldMergeMapAgg::Retract(const VariantType& accumulator, + const VariantType& input_field) const { + if (DataDefine::IsVariantNull(accumulator) || DataDefine::IsVariantNull(input_field)) { + return accumulator; + } + auto accumulator_map = DataDefine::GetVariantValue>(accumulator); + auto retract_map = DataDefine::GetVariantValue>(input_field); + if (retract_map->Size() == 0) { + return accumulator; + } + + std::vector entries; + entries.reserve(accumulator_map->Size()); + PAIMON_RETURN_NOT_OK(PutMap(accumulator_map, key_type_, value_type_, &entries)); + std::shared_ptr retract_keys = retract_map->KeyArray(); + for (int32_t i = 0; i < retract_map->Size(); ++i) { + PAIMON_ASSIGN_OR_RAISE(VariantType key, + FieldAggregateUtils::GetValue(*retract_keys, i, key_type_)); + PAIMON_ASSIGN_OR_RAISE(int32_t existing, FindKey(entries, key, key_type_)); + if (existing >= 0) { + entries.erase(entries.begin() + existing); + } + } + return MakeMap(std::move(entries), {accumulator_map->KeyArray()}, + {accumulator_map->ValueArray()}); +} + +} // namespace paimon diff --git a/src/paimon/core/mergetree/compact/aggregate/field_merge_map_agg.h b/src/paimon/core/mergetree/compact/aggregate/field_merge_map_agg.h new file mode 100644 index 000000000..cc8931fb1 --- /dev/null +++ b/src/paimon/core/mergetree/compact/aggregate/field_merge_map_agg.h @@ -0,0 +1,60 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * 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 + * + * 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 +#include + +#include "paimon/core/mergetree/compact/aggregate/field_aggregator.h" + +namespace paimon { + +/// Merges map fields and lets input values overwrite matching accumulator keys. +class FieldMergeMapAgg : public FieldAggregator { + public: + static constexpr char NAME[] = "merge_map"; + + /// Create a merge_map aggregator for a map field. + /// + /// @param field_type Type of the aggregated field. + /// @param field_name Name of the aggregated field. + /// @return A merge_map aggregator, or an error Status. + static Result> Create( + const std::shared_ptr& field_type, const std::string& field_name); + + VariantType Agg(const VariantType& accumulator, const VariantType& input_field) override; + Result AggResult(const VariantType& accumulator, + const VariantType& input_field) override; + Result Retract(const VariantType& accumulator, + const VariantType& input_field) const override; + + private: + FieldMergeMapAgg(const std::shared_ptr& field_type, + std::shared_ptr key_type, + std::shared_ptr value_type) + : FieldAggregator(NAME, field_type), + key_type_(std::move(key_type)), + value_type_(std::move(value_type)) {} + + Result AggImpl(const VariantType& accumulator, + const VariantType& input_field) const; + + std::shared_ptr key_type_; + std::shared_ptr value_type_; +}; + +} // namespace paimon diff --git a/src/paimon/core/mergetree/compact/aggregate/field_merge_map_agg_test.cpp b/src/paimon/core/mergetree/compact/aggregate/field_merge_map_agg_test.cpp new file mode 100644 index 000000000..9e950cb56 --- /dev/null +++ b/src/paimon/core/mergetree/compact/aggregate/field_merge_map_agg_test.cpp @@ -0,0 +1,87 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * 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 + * + * 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 "paimon/core/mergetree/compact/aggregate/field_merge_map_agg.h" + +#include +#include +#include + +#include "arrow/api.h" +#include "gtest/gtest.h" +#include "paimon/common/data/generic_array.h" +#include "paimon/common/data/generic_map.h" +#include "paimon/common/data/serializer/binary_serializer_utils.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { +namespace { + +VariantType IntMap(std::vector keys, std::vector values) { + std::shared_ptr key_array = std::make_shared(std::move(keys)); + std::shared_ptr value_array = std::make_shared(std::move(values)); + return VariantType(std::static_pointer_cast( + std::make_shared(std::move(key_array), std::move(value_array)))); +} + +int32_t FindValue(const VariantType& value, int32_t key) { + auto map = DataDefine::GetVariantValue>(value); + for (int32_t i = 0; i < map->Size(); ++i) { + if (map->KeyArray()->GetInt(i) == key) { + return map->ValueArray()->GetInt(i); + } + } + return -1; +} + +} // namespace + +TEST(FieldMergeMapAggTest, InputOverwritesAndRetractUsesKeys) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, + FieldMergeMapAgg::Create(arrow::map(arrow::int32(), arrow::int32()), "f")); + + ASSERT_OK_AND_ASSIGN( + VariantType merged, + agg->AggResult(IntMap({int32_t{1}, int32_t{2}}, {int32_t{10}, int32_t{20}}), + IntMap({int32_t{2}, int32_t{3}}, {int32_t{200}, int32_t{30}}))); + auto merged_map = DataDefine::GetVariantValue>(merged); + ASSERT_EQ(3, merged_map->Size()); + ASSERT_EQ(10, FindValue(merged, 1)); + ASSERT_EQ(200, FindValue(merged, 2)); + ASSERT_EQ(30, FindValue(merged, 3)); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr binary_merged, + BinarySerializerUtils::WriteBinaryMap( + merged_map, arrow::map(arrow::int32(), arrow::int32()), GetDefaultPool().get())); + ASSERT_EQ(3, binary_merged->Size()); + + ASSERT_OK_AND_ASSIGN(VariantType retracted, + agg->Retract(merged, IntMap({int32_t{2}}, {int32_t{-999}}))); + auto retracted_map = DataDefine::GetVariantValue>(retracted); + ASSERT_EQ(2, retracted_map->Size()); + ASSERT_EQ(-1, FindValue(retracted, 2)); +} + +TEST(FieldMergeMapAggTest, NullAndTypeValidation) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, + FieldMergeMapAgg::Create(arrow::map(arrow::int32(), arrow::int32()), "f")); + VariantType map = IntMap({int32_t{1}}, {int32_t{10}}); + ASSERT_OK_AND_ASSIGN(VariantType result, agg->AggResult(VariantType(NullType()), map)); + ASSERT_EQ(10, FindValue(result, 1)); + ASSERT_NOK(FieldMergeMapAgg::Create(arrow::int32(), "f")); +} + +} // namespace paimon::test diff --git a/src/paimon/core/mergetree/compact/aggregate/field_nested_update_agg.cpp b/src/paimon/core/mergetree/compact/aggregate/field_nested_update_agg.cpp new file mode 100644 index 000000000..716a9a0ee --- /dev/null +++ b/src/paimon/core/mergetree/compact/aggregate/field_nested_update_agg.cpp @@ -0,0 +1,370 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * 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 + * + * 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 "paimon/core/mergetree/compact/aggregate/field_nested_update_agg.h" + +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/util/checked_cast.h" +#include "fmt/format.h" +#include "paimon/common/data/generic_array.h" +#include "paimon/common/data/internal_row.h" +#include "paimon/common/types/data_field.h" +#include "paimon/common/utils/fields_comparator.h" +#include "paimon/core/mergetree/compact/aggregate/field_aggregate_utils.h" +#include "paimon/defs.h" +#include "paimon/status.h" + +namespace paimon { +namespace { + +std::string FieldOptionKey(const std::string& field_name, const char* option) { + return std::string(Options::FIELDS_PREFIX) + "." + field_name + "." + option; +} + +Result> ResolveFields(const std::shared_ptr& row_type, + const std::vector& names, + const std::string& option_name) { + std::vector fields; + fields.reserve(names.size()); + for (const std::string& name : names) { + int32_t index = row_type->GetFieldIndex(name); + if (index < 0) { + return Status::Invalid(fmt::format("Field '{}' configured by '{}' does not exist in {}", + name, option_name, row_type->ToString())); + } + if (std::find(fields.begin(), fields.end(), index) != fields.end()) { + return Status::Invalid( + fmt::format("Field '{}' is configured more than once by '{}'", name, option_name)); + } + fields.push_back(index); + } + return fields; +} + +std::shared_ptr MakeRows(std::vector> rows, + std::vector> holders) { + std::vector values; + values.reserve(rows.size()); + for (std::shared_ptr& row : rows) { + values.push_back(std::move(row)); + } + return std::make_shared(std::move(values), std::move(holders)); +} + +void AppendNonNullRows(const std::shared_ptr& array, int32_t row_fields, + int32_t limit, std::vector>* rows) { + int32_t added = 0; + for (int32_t i = 0; i < array->Size() && added < limit; ++i) { + if (!array->IsNullAt(i)) { + rows->push_back(array->GetRow(i, row_fields)); + ++added; + } + } +} + +} // namespace + +FieldNestedUpdateAgg::FieldNestedUpdateAgg(const std::shared_ptr& field_type, + std::shared_ptr row_type, + std::vector key_fields, + CoreOptions::NestedKeyNullStrategy null_strategy, + std::unique_ptr sequence_comparator, + int32_t count_limit) + : FieldAggregator(NAME, field_type), + row_type_(std::move(row_type)), + key_fields_(std::move(key_fields)), + null_strategy_(null_strategy), + sequence_comparator_(std::move(sequence_comparator)), + count_limit_(count_limit) {} + +FieldNestedUpdateAgg::~FieldNestedUpdateAgg() = default; + +Result> FieldNestedUpdateAgg::Create( + const std::shared_ptr& field_type, const CoreOptions& options, + const std::string& field_name) { + if (field_type->id() != arrow::Type::LIST) { + return Status::Invalid( + fmt::format("invalid field type {} for field '{}' of {}, supposed to be array", + field_type->ToString(), field_name, NAME)); + } + std::shared_ptr list_type = + arrow::internal::checked_pointer_cast(field_type); + if (list_type->value_type()->id() != arrow::Type::STRUCT) { + return Status::Invalid( + fmt::format("invalid field type {} for field '{}' of {}, supposed to be array", + field_type->ToString(), field_name, NAME)); + } + std::shared_ptr row_type = + arrow::internal::checked_pointer_cast(list_type->value_type()); + + PAIMON_ASSIGN_OR_RAISE(std::vector key_names, + options.FieldNestedUpdateAggNestedKey(field_name)); + PAIMON_ASSIGN_OR_RAISE(std::vector sequence_names, + options.FieldNestedUpdateAggNestedSequenceField(field_name)); + bool strategy_configured = + options.ToMap().count(FieldOptionKey(field_name, Options::NESTED_KEY_NULL_STRATEGY)) > 0; + if (key_names.empty() && strategy_configured) { + return Status::Invalid( + "Option 'fields..nested-key-null-strategy' requires " + "'fields..nested-key' to be configured."); + } + if (key_names.empty() && !sequence_names.empty()) { + return Status::Invalid( + "Option 'fields..nested-sequence-field' requires " + "'fields..nested-key' to be configured."); + } + + PAIMON_ASSIGN_OR_RAISE(std::vector key_fields, + ResolveFields(row_type, key_names, Options::NESTED_KEY)); + PAIMON_ASSIGN_OR_RAISE(std::vector sequence_fields, + ResolveFields(row_type, sequence_names, Options::NESTED_SEQUENCE_FIELD)); + std::unique_ptr sequence_comparator; + if (!sequence_fields.empty()) { + std::vector row_fields; + row_fields.reserve(row_type->num_fields()); + for (int32_t i = 0; i < row_type->num_fields(); ++i) { + row_fields.emplace_back(i, row_type->field(i)); + } + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr comparator, + FieldsComparator::Create(row_fields, sequence_fields, + /*is_ascending_order=*/true)); + sequence_comparator = std::move(comparator); + } + PAIMON_ASSIGN_OR_RAISE(CoreOptions::NestedKeyNullStrategy null_strategy, + options.FieldNestedUpdateAggNestedKeyNullStrategy(field_name)); + PAIMON_ASSIGN_OR_RAISE(int32_t count_limit, options.FieldNestedUpdateAggCountLimit(field_name)); + return std::unique_ptr( + new FieldNestedUpdateAgg(field_type, std::move(row_type), std::move(key_fields), + null_strategy, std::move(sequence_comparator), count_limit)); +} + +VariantType FieldNestedUpdateAgg::Agg(const VariantType& accumulator, + const VariantType& input_field) { + Result result = AggImpl(accumulator, input_field); + return result.ok() ? std::move(result).value() : accumulator; +} + +Result FieldNestedUpdateAgg::AggResult(const VariantType& accumulator, + const VariantType& input_field) { + return AggImpl(accumulator, input_field); +} + +Result FieldNestedUpdateAgg::AcceptKey(const InternalRow& row) const { + bool contains_null = false; + for (int32_t field : key_fields_) { + contains_null = contains_null || row.IsNullAt(field); + } + if (!contains_null || null_strategy_ == CoreOptions::NestedKeyNullStrategy::MERGE) { + return true; + } + if (null_strategy_ == CoreOptions::NestedKeyNullStrategy::IGNORE) { + return false; + } + return Status::Invalid("Nested key contains null values. Primary key fields must not be null."); +} + +Result FieldNestedUpdateAgg::KeysEqual(const InternalRow& lhs, const InternalRow& rhs) const { + for (int32_t field : key_fields_) { + PAIMON_ASSIGN_OR_RAISE( + VariantType lhs_value, + FieldAggregateUtils::GetValue(lhs, field, row_type_->field(field)->type())); + PAIMON_ASSIGN_OR_RAISE( + VariantType rhs_value, + FieldAggregateUtils::GetValue(rhs, field, row_type_->field(field)->type())); + PAIMON_ASSIGN_OR_RAISE( + bool equal, + FieldAggregateUtils::Equals(lhs_value, rhs_value, row_type_->field(field)->type())); + if (!equal) { + return false; + } + } + return true; +} + +Result FieldNestedUpdateAgg::RowsEqual(const InternalRow& lhs, const InternalRow& rhs) const { + for (int32_t field = 0; field < row_type_->num_fields(); ++field) { + PAIMON_ASSIGN_OR_RAISE( + VariantType lhs_value, + FieldAggregateUtils::GetValue(lhs, field, row_type_->field(field)->type())); + PAIMON_ASSIGN_OR_RAISE( + VariantType rhs_value, + FieldAggregateUtils::GetValue(rhs, field, row_type_->field(field)->type())); + PAIMON_ASSIGN_OR_RAISE( + bool equal, + FieldAggregateUtils::Equals(lhs_value, rhs_value, row_type_->field(field)->type())); + if (!equal) { + return false; + } + } + return true; +} + +Result FieldNestedUpdateAgg::AggImpl(const VariantType& accumulator, + const VariantType& input_field) const { + if (DataDefine::IsVariantNull(input_field)) { + return accumulator; + } + auto input = DataDefine::GetVariantValue>(input_field); + std::shared_ptr acc = + DataDefine::IsVariantNull(accumulator) + ? nullptr + : DataDefine::GetVariantValue>(accumulator); + + if (key_fields_.empty()) { + if (acc && acc->Size() >= count_limit_) { + return accumulator; + } + std::vector> rows; + if (acc) { + rows.reserve(acc->Size() + input->Size()); + AppendNonNullRows(acc, row_type_->num_fields(), acc->Size(), &rows); + } + int32_t remaining = acc ? count_limit_ - acc->Size() : count_limit_; + AppendNonNullRows(input, row_type_->num_fields(), remaining, &rows); + std::vector> holders; + if (acc) { + holders.push_back(acc); + } + holders.push_back(input); + return VariantType( + std::static_pointer_cast(MakeRows(std::move(rows), std::move(holders)))); + } + + std::vector> rows; + auto add_rows = [&](const std::shared_ptr& array, + bool limit_new_keys) -> Status { + if (!array) { + return Status::OK(); + } + for (int32_t i = 0; i < array->Size(); ++i) { + if (array->IsNullAt(i)) { + continue; + } + std::shared_ptr row = array->GetRow(i, row_type_->num_fields()); + PAIMON_ASSIGN_OR_RAISE(bool accept, AcceptKey(*row)); + if (!accept) { + continue; + } + int32_t existing = -1; + for (int32_t j = 0; j < static_cast(rows.size()); ++j) { + PAIMON_ASSIGN_OR_RAISE(bool equal, KeysEqual(*rows[j], *row)); + if (equal) { + existing = j; + break; + } + } + if (existing >= 0) { + if (!sequence_comparator_ || + sequence_comparator_->CompareTo(*row, *rows[existing]) >= 0) { + rows[existing] = std::move(row); + } + } else if (!limit_new_keys || static_cast(rows.size()) < count_limit_) { + rows.push_back(std::move(row)); + } + } + return Status::OK(); + }; + PAIMON_RETURN_NOT_OK(add_rows(acc, /*limit_new_keys=*/false)); + PAIMON_RETURN_NOT_OK(add_rows(input, /*limit_new_keys=*/true)); + std::vector> holders; + if (acc) { + holders.push_back(acc); + } + holders.push_back(input); + return VariantType( + std::static_pointer_cast(MakeRows(std::move(rows), std::move(holders)))); +} + +Result FieldNestedUpdateAgg::Retract(const VariantType& accumulator, + const VariantType& input_field) const { + if (DataDefine::IsVariantNull(accumulator) || DataDefine::IsVariantNull(input_field)) { + return accumulator; + } + auto acc = DataDefine::GetVariantValue>(accumulator); + auto retract = DataDefine::GetVariantValue>(input_field); + std::vector> rows; + + if (key_fields_.empty()) { + AppendNonNullRows(acc, row_type_->num_fields(), acc->Size(), &rows); + for (int32_t i = 0; i < retract->Size(); ++i) { + if (retract->IsNullAt(i)) { + continue; + } + std::shared_ptr retract_row = retract->GetRow(i, row_type_->num_fields()); + for (auto iter = rows.begin(); iter != rows.end();) { + PAIMON_ASSIGN_OR_RAISE(bool equal, RowsEqual(**iter, *retract_row)); + if (equal) { + iter = rows.erase(iter); + } else { + ++iter; + } + } + } + return VariantType(std::static_pointer_cast( + MakeRows(std::move(rows), std::vector>{acc, retract}))); + } + + for (int32_t i = 0; i < acc->Size(); ++i) { + if (acc->IsNullAt(i)) { + continue; + } + std::shared_ptr row = acc->GetRow(i, row_type_->num_fields()); + PAIMON_ASSIGN_OR_RAISE(bool accept, AcceptKey(*row)); + if (!accept) { + continue; + } + int32_t existing = -1; + for (int32_t j = 0; j < static_cast(rows.size()); ++j) { + PAIMON_ASSIGN_OR_RAISE(bool equal, KeysEqual(*rows[j], *row)); + if (equal) { + existing = j; + break; + } + } + if (existing >= 0) { + rows[existing] = std::move(row); + } else { + rows.push_back(std::move(row)); + } + } + + for (int32_t i = 0; i < retract->Size(); ++i) { + if (retract->IsNullAt(i)) { + continue; + } + std::shared_ptr retract_row = retract->GetRow(i, row_type_->num_fields()); + PAIMON_ASSIGN_OR_RAISE(bool accept, AcceptKey(*retract_row)); + if (!accept) { + continue; + } + for (auto iter = rows.begin(); iter != rows.end();) { + PAIMON_ASSIGN_OR_RAISE(bool equal, KeysEqual(**iter, *retract_row)); + if (equal) { + iter = rows.erase(iter); + } else { + ++iter; + } + } + } + return VariantType(std::static_pointer_cast( + MakeRows(std::move(rows), std::vector>{acc, retract}))); +} + +} // namespace paimon diff --git a/src/paimon/core/mergetree/compact/aggregate/field_nested_update_agg.h b/src/paimon/core/mergetree/compact/aggregate/field_nested_update_agg.h new file mode 100644 index 000000000..5223cc325 --- /dev/null +++ b/src/paimon/core/mergetree/compact/aggregate/field_nested_update_agg.h @@ -0,0 +1,78 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * 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 + * + * 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 +#include +#include + +#include "paimon/core/core_options.h" +#include "paimon/core/mergetree/compact/aggregate/field_aggregator.h" + +namespace arrow { +class StructType; +} // namespace arrow + +namespace paimon { + +class FieldsComparator; + +/// Upserts rows in an ARRAY using configured nested keys and sequence fields. +class FieldNestedUpdateAgg : public FieldAggregator { + public: + static constexpr char NAME[] = "nested_update"; + + ~FieldNestedUpdateAgg() override; + + /// Create a nested_update aggregator for an array-of-struct field. + /// + /// @param field_type Type of the aggregated field. + /// @param options Table options describing nested keys, sequences, and limits. + /// @param field_name Name of the aggregated field. + /// @return A nested_update aggregator, or an error Status. + static Result> Create( + const std::shared_ptr& field_type, const CoreOptions& options, + const std::string& field_name); + + VariantType Agg(const VariantType& accumulator, const VariantType& input_field) override; + Result AggResult(const VariantType& accumulator, + const VariantType& input_field) override; + Result Retract(const VariantType& accumulator, + const VariantType& input_field) const override; + + private: + FieldNestedUpdateAgg(const std::shared_ptr& field_type, + std::shared_ptr row_type, + std::vector key_fields, + CoreOptions::NestedKeyNullStrategy null_strategy, + std::unique_ptr sequence_comparator, + int32_t count_limit); + + Result AggImpl(const VariantType& accumulator, + const VariantType& input_field) const; + Result AcceptKey(const InternalRow& row) const; + Result KeysEqual(const InternalRow& lhs, const InternalRow& rhs) const; + Result RowsEqual(const InternalRow& lhs, const InternalRow& rhs) const; + + std::shared_ptr row_type_; + std::vector key_fields_; + CoreOptions::NestedKeyNullStrategy null_strategy_; + std::unique_ptr sequence_comparator_; + int32_t count_limit_; +}; + +} // namespace paimon diff --git a/src/paimon/core/mergetree/compact/aggregate/field_nested_update_agg_test.cpp b/src/paimon/core/mergetree/compact/aggregate/field_nested_update_agg_test.cpp new file mode 100644 index 000000000..30e8dd1c8 --- /dev/null +++ b/src/paimon/core/mergetree/compact/aggregate/field_nested_update_agg_test.cpp @@ -0,0 +1,190 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * 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 + * + * 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 "paimon/core/mergetree/compact/aggregate/field_nested_update_agg.h" + +#include +#include +#include +#include +#include + +#include "arrow/api.h" +#include "gtest/gtest.h" +#include "paimon/common/data/generic_array.h" +#include "paimon/common/data/generic_row.h" +#include "paimon/common/data/serializer/binary_serializer_utils.h" +#include "paimon/core/core_options.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { +namespace { + +std::shared_ptr NestedType() { + return arrow::list( + arrow::struct_({arrow::field("id", arrow::int32()), arrow::field("seq", arrow::int32()), + arrow::field("value", arrow::int32())})); +} + +std::shared_ptr Row(VariantType id, int32_t sequence, int32_t value) { + std::shared_ptr row = std::make_shared(3); + row->SetField(0, id); + row->SetField(1, sequence); + row->SetField(2, value); + return row; +} + +VariantType Rows(std::vector rows) { + return VariantType( + std::static_pointer_cast(std::make_shared(std::move(rows)))); +} + +std::shared_ptr GetRows(const VariantType& value) { + return DataDefine::GetVariantValue>(value); +} + +std::shared_ptr FindRow(const VariantType& value, int32_t id) { + std::shared_ptr rows = GetRows(value); + for (int32_t i = 0; i < rows->Size(); ++i) { + if (rows->IsNullAt(i)) { + continue; + } + std::shared_ptr row = rows->GetRow(i, 3); + if (!row->IsNullAt(0) && row->GetInt(0) == id) { + return row; + } + } + return nullptr; +} + +Result> MakeAgg( + const std::map& options_map) { + PAIMON_ASSIGN_OR_RAISE(CoreOptions options, CoreOptions::FromMap(options_map)); + return FieldNestedUpdateAgg::Create(NestedType(), options, "f"); +} + +} // namespace + +TEST(FieldNestedUpdateAggTest, UpsertsByKeySequenceAndCountLimit) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, + MakeAgg({{"fields.f.nested-key", "id"}, + {"fields.f.nested-sequence-field", "seq"}, + {"fields.f.count-limit", "2"}})); + + VariantType accumulator = Rows({Row(int32_t{1}, 1, 10), Row(int32_t{2}, 1, 20)}); + VariantType input = + Rows({Row(int32_t{1}, 0, 100), Row(int32_t{1}, 2, 200), Row(int32_t{3}, 3, 300)}); + ASSERT_OK_AND_ASSIGN(VariantType result, agg->AggResult(accumulator, input)); + + ASSERT_EQ(2, GetRows(result)->Size()); + ASSERT_EQ(200, FindRow(result, 1)->GetInt(2)); + ASSERT_EQ(20, FindRow(result, 2)->GetInt(2)); + ASSERT_FALSE(FindRow(result, 3)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr binary_result, + BinarySerializerUtils::WriteBinaryArray(GetRows(result), NestedType(), + GetDefaultPool().get())); + ASSERT_EQ(2, binary_result->Size()); + + ASSERT_OK_AND_ASSIGN(VariantType retracted, + agg->Retract(result, Rows({Row(int32_t{1}, 999, -1)}))); + ASSERT_EQ(1, GetRows(retracted)->Size()); + ASSERT_FALSE(FindRow(retracted, 1)); +} + +TEST(FieldNestedUpdateAggTest, AppendsNonNullRowsUpToLimit) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, + MakeAgg({{"fields.f.count-limit", "2"}})); + ASSERT_OK_AND_ASSIGN(VariantType result, + agg->AggResult(VariantType(NullType()), + Rows({VariantType(NullType()), Row(int32_t{1}, 1, 10), + Row(int32_t{2}, 1, 20), Row(int32_t{3}, 1, 30)}))); + ASSERT_EQ(2, GetRows(result)->Size()); + ASSERT_TRUE(FindRow(result, 1)); + ASSERT_TRUE(FindRow(result, 2)); +} + +// count limit is measured against the raw element count, so null elements consume the limit even +// though they are dropped from the result +TEST(FieldNestedUpdateAggTest, CountLimitCountsNullElementsOfAccumulator) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr full_agg, + MakeAgg({{"fields.f.count-limit", "3"}})); + VariantType full = Rows({VariantType(NullType()), Row(int32_t{1}, 1, 10), + VariantType(NullType()), Row(int32_t{2}, 1, 20)}); + ASSERT_OK_AND_ASSIGN(VariantType unchanged, + full_agg->AggResult(full, Rows({Row(int32_t{3}, 1, 30)}))); + ASSERT_EQ(4, GetRows(unchanged)->Size()); + ASSERT_FALSE(FindRow(unchanged, 3)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, + MakeAgg({{"fields.f.count-limit", "4"}})); + VariantType accumulator = Rows({VariantType(NullType()), Row(int32_t{1}, 1, 10)}); + ASSERT_OK_AND_ASSIGN( + VariantType result, + agg->AggResult(accumulator, Rows({Row(int32_t{2}, 1, 20), Row(int32_t{3}, 1, 30), + Row(int32_t{4}, 1, 40)}))); + ASSERT_EQ(3, GetRows(result)->Size()); + ASSERT_TRUE(FindRow(result, 1)); + ASSERT_TRUE(FindRow(result, 2)); + ASSERT_TRUE(FindRow(result, 3)); + ASSERT_FALSE(FindRow(result, 4)); +} + +TEST(FieldNestedUpdateAggTest, AppliesNullKeyStrategies) { + ASSERT_OK_AND_ASSIGN( + std::unique_ptr ignore_agg, + MakeAgg({{"fields.f.nested-key", "id"}, {"fields.f.nested-key-null-strategy", "ignore"}})); + ASSERT_OK_AND_ASSIGN(VariantType ignored, + ignore_agg->AggResult(VariantType(NullType()), + Rows({Row(VariantType(NullType()), 1, 10)}))); + ASSERT_EQ(0, GetRows(ignored)->Size()); + + ASSERT_OK_AND_ASSIGN(VariantType normalized, + ignore_agg->Retract(Rows({Row(VariantType(NullType()), 1, 10), + Row(int32_t{1}, 1, 10), Row(int32_t{1}, 2, 20)}), + Rows({}))); + ASSERT_EQ(1, GetRows(normalized)->Size()); + ASSERT_EQ(20, FindRow(normalized, 1)->GetInt(2)); + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr error_agg, + MakeAgg({{"fields.f.nested-key", "id"}, {"fields.f.nested-key-null-strategy", "error"}})); + ASSERT_NOK( + error_agg->AggResult(VariantType(NullType()), Rows({Row(VariantType(NullType()), 1, 10)}))); +} + +TEST(FieldNestedUpdateAggTest, ValidatesTypeAndOptionDependencies) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({})); + ASSERT_NOK(FieldNestedUpdateAgg::Create(arrow::list(arrow::int32()), options, "f")); + + ASSERT_OK_AND_ASSIGN(CoreOptions strategy_without_key, + CoreOptions::FromMap({{"fields.f.nested-key-null-strategy", "ignore"}})); + ASSERT_NOK(FieldNestedUpdateAgg::Create(NestedType(), strategy_without_key, "f")); + + ASSERT_OK_AND_ASSIGN(CoreOptions sequence_without_key, + CoreOptions::FromMap({{"fields.f.nested-sequence-field", "seq"}})); + ASSERT_NOK(FieldNestedUpdateAgg::Create(NestedType(), sequence_without_key, "f")); + + ASSERT_OK_AND_ASSIGN(CoreOptions invalid_strategy, + CoreOptions::FromMap({{"fields.f.nested-key", "id"}, + {"fields.f.nested-key-null-strategy", "invalid"}})); + ASSERT_NOK(FieldNestedUpdateAgg::Create(NestedType(), invalid_strategy, "f")); + + ASSERT_OK_AND_ASSIGN(CoreOptions negative_limit, + CoreOptions::FromMap({{"fields.f.count-limit", "-1"}})); + ASSERT_NOK(FieldNestedUpdateAgg::Create(NestedType(), negative_limit, "f")); +} + +} // namespace paimon::test diff --git a/src/paimon/core/mergetree/compact/aggregate/field_roaring_bitmap64_agg.cpp b/src/paimon/core/mergetree/compact/aggregate/field_roaring_bitmap64_agg.cpp new file mode 100644 index 000000000..f5e681d73 --- /dev/null +++ b/src/paimon/core/mergetree/compact/aggregate/field_roaring_bitmap64_agg.cpp @@ -0,0 +1,62 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * 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 + * + * 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 "paimon/core/mergetree/compact/aggregate/field_roaring_bitmap64_agg.h" + +#include "arrow/api.h" +#include "fmt/format.h" +#include "paimon/core/mergetree/compact/aggregate/field_aggregate_utils.h" +#include "paimon/memory/bytes.h" +#include "paimon/status.h" +#include "paimon/utils/roaring_bitmap64.h" + +namespace paimon { + +Result> FieldRoaringBitmap64Agg::Create( + const std::shared_ptr& field_type, const std::string& field_name) { + if (field_type->id() != arrow::Type::BINARY) { + return Status::Invalid( + fmt::format("invalid field type {} for field '{}' of {}, supposed to be binary", + field_type->ToString(), field_name, NAME)); + } + return std::unique_ptr(new FieldRoaringBitmap64Agg(field_type)); +} + +VariantType FieldRoaringBitmap64Agg::Agg(const VariantType& accumulator, + const VariantType& input_field) { + Result result = AggResult(accumulator, input_field); + return result.ok() ? std::move(result).value() : accumulator; +} + +Result FieldRoaringBitmap64Agg::AggResult(const VariantType& accumulator, + const VariantType& input_field) { + bool accumulator_null = DataDefine::IsVariantNull(accumulator); + bool input_null = DataDefine::IsVariantNull(input_field); + if (accumulator_null || input_null) { + return accumulator_null ? input_field : FieldAggregateUtils::OwnedBinary(accumulator); + } + std::string_view accumulator_bytes = DataDefine::GetStringView(accumulator); + std::string_view input_bytes = DataDefine::GetStringView(input_field); + RoaringBitmap64 accumulator_bitmap; + RoaringBitmap64 input_bitmap; + PAIMON_RETURN_NOT_OK( + accumulator_bitmap.Deserialize(accumulator_bytes.data(), accumulator_bytes.size())); + PAIMON_RETURN_NOT_OK(input_bitmap.Deserialize(input_bytes.data(), input_bytes.size())); + accumulator_bitmap |= input_bitmap; + return VariantType(std::shared_ptr(accumulator_bitmap.Serialize(/*pool=*/nullptr))); +} + +} // namespace paimon diff --git a/src/paimon/core/mergetree/compact/aggregate/field_roaring_bitmap64_agg.h b/src/paimon/core/mergetree/compact/aggregate/field_roaring_bitmap64_agg.h new file mode 100644 index 000000000..efd1db527 --- /dev/null +++ b/src/paimon/core/mergetree/compact/aggregate/field_roaring_bitmap64_agg.h @@ -0,0 +1,48 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * 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 + * + * 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 +#include + +#include "paimon/core/mergetree/compact/aggregate/field_aggregator.h" + +namespace paimon { + +/// Unions serialized 64-bit Roaring Bitmap fields. +class FieldRoaringBitmap64Agg : public FieldAggregator { + public: + static constexpr char NAME[] = "rbm64"; + + /// Create an rbm64 aggregator for a binary field. + /// + /// @param field_type Type of the aggregated field. + /// @param field_name Name of the aggregated field. + /// @return An rbm64 aggregator, or an error Status. + static Result> Create( + const std::shared_ptr& field_type, const std::string& field_name); + + VariantType Agg(const VariantType& accumulator, const VariantType& input_field) override; + Result AggResult(const VariantType& accumulator, + const VariantType& input_field) override; + + private: + explicit FieldRoaringBitmap64Agg(const std::shared_ptr& field_type) + : FieldAggregator(NAME, field_type) {} +}; + +} // namespace paimon diff --git a/src/paimon/core/mergetree/compact/aggregate/field_roaring_bitmap64_agg_test.cpp b/src/paimon/core/mergetree/compact/aggregate/field_roaring_bitmap64_agg_test.cpp new file mode 100644 index 000000000..5093e8364 --- /dev/null +++ b/src/paimon/core/mergetree/compact/aggregate/field_roaring_bitmap64_agg_test.cpp @@ -0,0 +1,62 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * 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 + * + * 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 "paimon/core/mergetree/compact/aggregate/field_roaring_bitmap64_agg.h" + +#include + +#include "arrow/api.h" +#include "gtest/gtest.h" +#include "paimon/memory/bytes.h" +#include "paimon/testing/utils/testharness.h" +#include "paimon/utils/roaring_bitmap64.h" + +namespace paimon::test { +namespace { + +VariantType Bitmap(std::initializer_list values) { + RoaringBitmap64 bitmap; + for (int64_t value : values) { + bitmap.Add(value); + } + return VariantType(std::shared_ptr(bitmap.Serialize(/*pool=*/nullptr))); +} + +} // namespace + +TEST(FieldRoaringBitmap64AggTest, UnionsSerializedBitmaps) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, + FieldRoaringBitmap64Agg::Create(arrow::binary(), "f")); + ASSERT_OK_AND_ASSIGN(VariantType result, + agg->AggResult(Bitmap({1, 3}), Bitmap({2, 3, int64_t{1} << 40}))); + std::string_view bytes = DataDefine::GetStringView(result); + RoaringBitmap64 bitmap; + ASSERT_OK(bitmap.Deserialize(bytes.data(), bytes.size())); + ASSERT_EQ(4, bitmap.Cardinality()); + ASSERT_TRUE(bitmap.Contains(1)); + ASSERT_TRUE(bitmap.Contains(2)); + ASSERT_TRUE(bitmap.Contains(3)); + ASSERT_TRUE(bitmap.Contains(int64_t{1} << 40)); +} + +TEST(FieldRoaringBitmap64AggTest, ReportsInvalidBytesAndType) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, + FieldRoaringBitmap64Agg::Create(arrow::binary(), "f")); + ASSERT_NOK(agg->AggResult(VariantType(std::string_view("bad")), Bitmap({1}))); + ASSERT_NOK(FieldRoaringBitmap64Agg::Create(arrow::utf8(), "f")); +} + +} // namespace paimon::test diff --git a/src/paimon/core/mergetree/compact/aggregate/field_sketch_agg.cpp b/src/paimon/core/mergetree/compact/aggregate/field_sketch_agg.cpp new file mode 100644 index 000000000..a77f917ca --- /dev/null +++ b/src/paimon/core/mergetree/compact/aggregate/field_sketch_agg.cpp @@ -0,0 +1,136 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * 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 + * + * 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 "paimon/core/mergetree/compact/aggregate/field_sketch_agg.h" + +#include +#include +#include +#include + +#include "DataSketches/hll.hpp" +#include "DataSketches/theta_sketch.hpp" +#include "DataSketches/theta_union.hpp" +#include "arrow/api.h" +#include "fmt/format.h" +#include "paimon/core/mergetree/compact/aggregate/field_aggregate_utils.h" +#include "paimon/memory/bytes.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/status.h" + +namespace paimon { +namespace { + +template +std::shared_ptr CopyBytes(const std::vector& serialized) { + pooled_unique_ptr result = + Bytes::AllocateBytes(serialized.size() * sizeof(T), GetDefaultPool().get()); + if (!serialized.empty()) { + std::memcpy(result->data(), serialized.data(), result->size()); + } + return std::shared_ptr(std::move(result)); +} + +Status ValidateSketchType(const std::shared_ptr& field_type, + const std::string& field_name, const char* name) { + if (field_type->id() != arrow::Type::BINARY) { + return Status::Invalid( + fmt::format("invalid field type {} for field '{}' of {}, supposed to be binary", + field_type->ToString(), field_name, name)); + } + return Status::OK(); +} + +} // namespace + +Result> FieldHllSketchAgg::Create( + const std::shared_ptr& field_type, const std::string& field_name) { + PAIMON_RETURN_NOT_OK(ValidateSketchType(field_type, field_name, NAME)); + return std::unique_ptr(new FieldHllSketchAgg(field_type)); +} + +VariantType FieldHllSketchAgg::Agg(const VariantType& accumulator, const VariantType& input_field) { + Result result = AggResult(accumulator, input_field); + return result.ok() ? std::move(result).value() : accumulator; +} + +Result FieldHllSketchAgg::AggResult(const VariantType& accumulator, + const VariantType& input_field) { + bool accumulator_null = DataDefine::IsVariantNull(accumulator); + bool input_null = DataDefine::IsVariantNull(input_field); + if (accumulator_null || input_null) { + return accumulator_null ? input_field : FieldAggregateUtils::OwnedBinary(accumulator); + } + std::string_view accumulator_bytes = DataDefine::GetStringView(accumulator); + std::string_view input_bytes = DataDefine::GetStringView(input_field); + try { + datasketches::hll_sketch accumulator_sketch = datasketches::hll_sketch::deserialize( + accumulator_bytes.data(), accumulator_bytes.size()); + datasketches::hll_sketch input_sketch = + datasketches::hll_sketch::deserialize(input_bytes.data(), input_bytes.size()); + datasketches::hll_union sketch_union(input_sketch.get_lg_config_k()); + sketch_union.update(input_sketch); + sketch_union.update(accumulator_sketch); + datasketches::hll_sketch result = sketch_union.get_result(datasketches::HLL_4); + return VariantType(CopyBytes(result.serialize_compact())); + } catch (const std::exception& exception) { + return Status::Invalid( + fmt::format("Unable to deserialize or union HLL sketch: {}", exception.what())); + } catch (...) { + return Status::Invalid("Unable to deserialize or union HLL sketch"); + } +} + +Result> FieldThetaSketchAgg::Create( + const std::shared_ptr& field_type, const std::string& field_name) { + PAIMON_RETURN_NOT_OK(ValidateSketchType(field_type, field_name, NAME)); + return std::unique_ptr(new FieldThetaSketchAgg(field_type)); +} + +VariantType FieldThetaSketchAgg::Agg(const VariantType& accumulator, + const VariantType& input_field) { + Result result = AggResult(accumulator, input_field); + return result.ok() ? std::move(result).value() : accumulator; +} + +Result FieldThetaSketchAgg::AggResult(const VariantType& accumulator, + const VariantType& input_field) { + bool accumulator_null = DataDefine::IsVariantNull(accumulator); + bool input_null = DataDefine::IsVariantNull(input_field); + if (accumulator_null || input_null) { + return accumulator_null ? input_field : FieldAggregateUtils::OwnedBinary(accumulator); + } + std::string_view accumulator_bytes = DataDefine::GetStringView(accumulator); + std::string_view input_bytes = DataDefine::GetStringView(input_field); + try { + datasketches::compact_theta_sketch accumulator_sketch = + datasketches::compact_theta_sketch::deserialize(accumulator_bytes.data(), + accumulator_bytes.size()); + datasketches::compact_theta_sketch input_sketch = + datasketches::compact_theta_sketch::deserialize(input_bytes.data(), input_bytes.size()); + datasketches::theta_union sketch_union = datasketches::theta_union::builder().build(); + sketch_union.update(accumulator_sketch); + sketch_union.update(input_sketch); + return VariantType(CopyBytes(sketch_union.get_result(/*ordered=*/true).serialize())); + } catch (const std::exception& exception) { + return Status::Invalid( + fmt::format("Unable to deserialize or union theta sketch: {}", exception.what())); + } catch (...) { + return Status::Invalid("Unable to deserialize or union theta sketch"); + } +} + +} // namespace paimon diff --git a/src/paimon/core/mergetree/compact/aggregate/field_sketch_agg.h b/src/paimon/core/mergetree/compact/aggregate/field_sketch_agg.h new file mode 100644 index 000000000..a0dfa1eca --- /dev/null +++ b/src/paimon/core/mergetree/compact/aggregate/field_sketch_agg.h @@ -0,0 +1,70 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * 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 + * + * 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 +#include + +#include "paimon/core/mergetree/compact/aggregate/field_aggregator.h" + +namespace paimon { + +/// Unions serialized HyperLogLog sketch fields. +class FieldHllSketchAgg : public FieldAggregator { + public: + static constexpr char NAME[] = "hll_sketch"; + + /// Create an hll_sketch aggregator for a binary field. + /// + /// @param field_type Type of the aggregated field. + /// @param field_name Name of the aggregated field. + /// @return An hll_sketch aggregator, or an error Status. + static Result> Create( + const std::shared_ptr& field_type, const std::string& field_name); + + VariantType Agg(const VariantType& accumulator, const VariantType& input_field) override; + Result AggResult(const VariantType& accumulator, + const VariantType& input_field) override; + + private: + explicit FieldHllSketchAgg(const std::shared_ptr& field_type) + : FieldAggregator(NAME, field_type) {} +}; + +/// Unions serialized Theta sketch fields. +class FieldThetaSketchAgg : public FieldAggregator { + public: + static constexpr char NAME[] = "theta_sketch"; + + /// Create a theta_sketch aggregator for a binary field. + /// + /// @param field_type Type of the aggregated field. + /// @param field_name Name of the aggregated field. + /// @return A theta_sketch aggregator, or an error Status. + static Result> Create( + const std::shared_ptr& field_type, const std::string& field_name); + + VariantType Agg(const VariantType& accumulator, const VariantType& input_field) override; + Result AggResult(const VariantType& accumulator, + const VariantType& input_field) override; + + private: + explicit FieldThetaSketchAgg(const std::shared_ptr& field_type) + : FieldAggregator(NAME, field_type) {} +}; + +} // namespace paimon diff --git a/src/paimon/core/mergetree/compact/aggregate/field_sketch_agg_test.cpp b/src/paimon/core/mergetree/compact/aggregate/field_sketch_agg_test.cpp new file mode 100644 index 000000000..a67b79acd --- /dev/null +++ b/src/paimon/core/mergetree/compact/aggregate/field_sketch_agg_test.cpp @@ -0,0 +1,186 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * 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 + * + * 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 "paimon/core/mergetree/compact/aggregate/field_sketch_agg.h" + +#include +#include +#include +#include +#include + +#include "DataSketches/hll.hpp" +#include "DataSketches/theta_sketch.hpp" +#include "arrow/api.h" +#include "gtest/gtest.h" +#include "paimon/common/types/row_kind.h" +#include "paimon/core/core_options.h" +#include "paimon/core/key_value.h" +#include "paimon/core/mergetree/compact/aggregate/aggregate_merge_function.h" +#include "paimon/defs.h" +#include "paimon/memory/bytes.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/testing/utils/binary_row_generator.h" +#include "paimon/testing/utils/testharness.h" +#include "paimon/utils/roaring_bitmap64.h" + +namespace paimon::test { +namespace { + +template +VariantType Serialized(const std::vector& data) { + std::shared_ptr bytes = + Bytes::AllocateBytes(data.size() * sizeof(T), GetDefaultPool().get()); + if (!data.empty()) { + std::memcpy(bytes->data(), data.data(), bytes->size()); + } + return VariantType(std::move(bytes)); +} + +VariantType Hll(std::initializer_list values) { + datasketches::hll_sketch sketch(12, datasketches::HLL_4); + for (int32_t value : values) { + sketch.update(value); + } + return Serialized(sketch.serialize_compact()); +} + +VariantType Theta(std::initializer_list values) { + datasketches::update_theta_sketch sketch = datasketches::update_theta_sketch::builder().build(); + for (int32_t value : values) { + sketch.update(value); + } + return Serialized(sketch.compact(/*ordered=*/true).serialize()); +} + +std::shared_ptr HllBytes(std::initializer_list values) { + return DataDefine::GetVariantValue>(Hll(values)); +} + +std::shared_ptr ThetaBytes(std::initializer_list values) { + return DataDefine::GetVariantValue>(Theta(values)); +} + +std::shared_ptr BitmapBytes(std::initializer_list values) { + RoaringBitmap64 bitmap; + for (int64_t value : values) { + bitmap.Add(value); + } + return std::shared_ptr(bitmap.Serialize(/*pool=*/nullptr)); +} + +// reuse freed heap blocks so a row still pointing into released memory yields corrupted bytes +std::vector> ScribbleFreedMemory() { + std::vector> blocks; + for (int32_t i = 0; i < 64; ++i) { + pooled_unique_ptr block = Bytes::AllocateBytes(4096, GetDefaultPool().get()); + std::memset(block->data(), 0xAB, block->size()); + blocks.push_back(std::move(block)); + } + return blocks; +} + +} // namespace + +TEST(BinaryAggMergeFunctionTest, OwnedAccumulatorSurvivesNullInput) { + arrow::FieldVector fields = { + arrow::field("k0", arrow::int32()), arrow::field("rbm", arrow::binary()), + arrow::field("hll", arrow::binary()), arrow::field("theta", arrow::binary())}; + ASSERT_OK_AND_ASSIGN( + CoreOptions options, + CoreOptions::FromMap({{"fields.rbm.aggregate-function", "rbm64"}, + {"fields.hll.aggregate-function", "hll_sketch"}, + {"fields.theta.aggregate-function", "theta_sketch"}})); + ASSERT_OK_AND_ASSIGN(std::unique_ptr merge_func, + AggregateMergeFunction::Create(arrow::schema(fields), + /*primary_keys=*/{"k0"}, options)); + + MemoryPool* pool = GetDefaultPool().get(); + ASSERT_OK(merge_func->Add(KeyValue( + RowKind::Insert(), /*sequence_number=*/0, /*level=*/0, + BinaryRowGenerator::GenerateRowPtr({10}, pool), + BinaryRowGenerator::GenerateRowPtr( + {10, BitmapBytes({1, 3}), HllBytes({1, 2, 3}), ThetaBytes({1, 2, 3})}, pool)))); + // both sides non-null, so each aggregator now owns a freshly allocated buffer in the row + ASSERT_OK(merge_func->Add(KeyValue( + RowKind::Insert(), /*sequence_number=*/1, /*level=*/0, + BinaryRowGenerator::GenerateRowPtr({10}, pool), + BinaryRowGenerator::GenerateRowPtr( + {10, BitmapBytes({2, 3}), HllBytes({3, 4, 5}), ThetaBytes({3, 4, 5})}, pool)))); + // input side is null, so the accumulator is passed through and written back into that field + ASSERT_OK(merge_func->Add(KeyValue( + RowKind::Insert(), /*sequence_number=*/2, /*level=*/0, + BinaryRowGenerator::GenerateRowPtr({10}, pool), + BinaryRowGenerator::GenerateRowPtr({10, NullType(), NullType(), NullType()}, pool)))); + + ASSERT_OK_AND_ASSIGN(std::optional result, merge_func->GetResult()); + ASSERT_TRUE(result.has_value()); + std::vector> scribbled = ScribbleFreedMemory(); + + std::string_view rbm_bytes = result->value->GetStringView(1); + RoaringBitmap64 bitmap; + ASSERT_OK(bitmap.Deserialize(rbm_bytes.data(), rbm_bytes.size())); + ASSERT_EQ(3, bitmap.Cardinality()); + + std::string_view hll_bytes = result->value->GetStringView(2); + ASSERT_NEAR( + 5.0, + datasketches::hll_sketch::deserialize(hll_bytes.data(), hll_bytes.size()).get_estimate(), + 0.1); + + std::string_view theta_bytes = result->value->GetStringView(3); + ASSERT_DOUBLE_EQ( + 5.0, datasketches::compact_theta_sketch::deserialize(theta_bytes.data(), theta_bytes.size()) + .get_estimate()); +} + +TEST(FieldSketchAggTest, UnionsHllSketchesAsCompactHll4) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, + FieldHllSketchAgg::Create(arrow::binary(), "f")); + ASSERT_OK_AND_ASSIGN(VariantType result, agg->AggResult(Hll({1, 2, 3}), Hll({3, 4, 5}))); + std::string_view bytes = DataDefine::GetStringView(result); + datasketches::hll_sketch sketch = + datasketches::hll_sketch::deserialize(bytes.data(), bytes.size()); + ASSERT_EQ(datasketches::HLL_4, sketch.get_target_type()); + ASSERT_NEAR(5.0, sketch.get_estimate(), 0.1); +} + +TEST(FieldSketchAggTest, UnionsOrderedThetaSketches) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, + FieldThetaSketchAgg::Create(arrow::binary(), "f")); + ASSERT_OK_AND_ASSIGN(VariantType result, agg->AggResult(Theta({1, 2, 3}), Theta({3, 4, 5}))); + std::string_view bytes = DataDefine::GetStringView(result); + datasketches::compact_theta_sketch sketch = + datasketches::compact_theta_sketch::deserialize(bytes.data(), bytes.size()); + ASSERT_TRUE(sketch.is_ordered()); + ASSERT_EQ(5, sketch.get_num_retained()); + ASSERT_DOUBLE_EQ(5.0, sketch.get_estimate()); +} + +TEST(FieldSketchAggTest, ReportsInvalidBytesAndTypes) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr hll_agg, + FieldHllSketchAgg::Create(arrow::binary(), "f")); + ASSERT_OK_AND_ASSIGN(std::unique_ptr theta_agg, + FieldThetaSketchAgg::Create(arrow::binary(), "f")); + VariantType invalid = VariantType(std::string_view("bad")); + ASSERT_NOK(hll_agg->AggResult(invalid, Hll({1}))); + ASSERT_NOK(hll_agg->AggReversedResult(invalid, Hll({1}))); + ASSERT_NOK(theta_agg->AggResult(invalid, Theta({1}))); + ASSERT_NOK(FieldHllSketchAgg::Create(arrow::int32(), "f")); + ASSERT_NOK(FieldThetaSketchAgg::Create(arrow::int32(), "f")); +} + +} // namespace paimon::test diff --git a/src/paimon/core/mergetree/compact/partial_update_merge_function.cpp b/src/paimon/core/mergetree/compact/partial_update_merge_function.cpp index 3ad9c3b28..a282ebbfc 100644 --- a/src/paimon/core/mergetree/compact/partial_update_merge_function.cpp +++ b/src/paimon/core/mergetree/compact/partial_update_merge_function.cpp @@ -333,16 +333,16 @@ Status PartialUpdateMergeFunction::Add(KeyValue&& moved_kv) { } last_seq_num_ = kv.sequence_number; if (field_comparators_.empty()) { - UpdateNonNullFields(std::move(kv)); + PAIMON_RETURN_NOT_OK(UpdateNonNullFields(std::move(kv))); } else { - UpdateWithSequenceGroup(std::move(kv)); + PAIMON_RETURN_NOT_OK(UpdateWithSequenceGroup(std::move(kv))); } meet_insert_ = true; not_null_column_filled_ = true; return Status::OK(); } -void PartialUpdateMergeFunction::UpdateNonNullFields(KeyValue&& kv) { +Status PartialUpdateMergeFunction::UpdateNonNullFields(KeyValue&& kv) { for (size_t i = 0; i < getters_.size(); ++i) { VariantType field = getters_[i](*(kv.value)); if (!DataDefine::IsVariantNull(field)) { @@ -350,9 +350,10 @@ void PartialUpdateMergeFunction::UpdateNonNullFields(KeyValue&& kv) { } } row_->AddDataHolder(std::move(kv.value)); + return Status::OK(); } -void PartialUpdateMergeFunction::UpdateWithSequenceGroup(KeyValue&& kv) { +Status PartialUpdateMergeFunction::UpdateWithSequenceGroup(KeyValue&& kv) { for (size_t i = 0; i < getters_.size(); ++i) { VariantType field = getters_[i](*(kv.value)); VariantType accumulator = getters_[i](*row_); @@ -362,7 +363,8 @@ void PartialUpdateMergeFunction::UpdateWithSequenceGroup(KeyValue&& kv) { (agg_iter == field_aggregators_.end() ? nullptr : agg_iter->second.get()); if (comp_iter == field_comparators_.end()) { if (agg) { - row_->SetField(i, agg->Agg(accumulator, field)); + PAIMON_ASSIGN_OR_RAISE(VariantType result, agg->AggResult(accumulator, field)); + row_->SetField(i, result); } else if (!DataDefine::IsVariantNull(field)) { row_->SetField(i, field); } @@ -383,13 +385,21 @@ void PartialUpdateMergeFunction::UpdateWithSequenceGroup(KeyValue&& kv) { } continue; } - row_->SetField(i, agg ? agg->Agg(accumulator, field) : field); + if (agg) { + PAIMON_ASSIGN_OR_RAISE(VariantType result, agg->AggResult(accumulator, field)); + row_->SetField(i, result); + } else { + row_->SetField(i, field); + } } else if (agg) { - row_->SetField(i, agg->AggReversed(accumulator, field)); + PAIMON_ASSIGN_OR_RAISE(VariantType result, + agg->AggReversedResult(accumulator, field)); + row_->SetField(i, result); } } } row_->AddDataHolder(std::move(kv.value)); + return Status::OK(); } Status PartialUpdateMergeFunction::RetractWithSequenceGroup(KeyValue&& kv) { diff --git a/src/paimon/core/mergetree/compact/partial_update_merge_function.h b/src/paimon/core/mergetree/compact/partial_update_merge_function.h index 5db121717..ce27b4c9b 100644 --- a/src/paimon/core/mergetree/compact/partial_update_merge_function.h +++ b/src/paimon/core/mergetree/compact/partial_update_merge_function.h @@ -106,9 +106,9 @@ class PartialUpdateMergeFunction : public MergeFunction { /// Initialize row_ with all field values and transfer data ownership to row_. void InitRowAndHoldData(std::unique_ptr&& value); - void UpdateNonNullFields(KeyValue&& kv); + Status UpdateNonNullFields(KeyValue&& kv); - void UpdateWithSequenceGroup(KeyValue&& kv); + Status UpdateWithSequenceGroup(KeyValue&& kv); Status RetractWithSequenceGroup(KeyValue&& kv); diff --git a/third_party/versions.txt b/third_party/versions.txt index d27b9e818..263a3a21f 100644 --- a/third_party/versions.txt +++ b/third_party/versions.txt @@ -113,6 +113,10 @@ PAIMON_RAPIDJSON_BUILD_VERSION=232389d4f1012dddec4ef84861face2d2ba85709 PAIMON_RAPIDJSON_BUILD_SHA256_CHECKSUM=b9290a9a6d444c8e049bd589ab804e0ccf2b05dc5984a19ed5ae75d090064806 PAIMON_RAPIDJSON_PKG_NAME=rapidjson-${PAIMON_RAPIDJSON_BUILD_VERSION}.tar.gz +PAIMON_DATASKETCHES_BUILD_VERSION=5.2.0 +PAIMON_DATASKETCHES_BUILD_SHA256_CHECKSUM=63e6bda660ee9730cc39746de87fd983118c28149e32ff9389fac98b27010a53 +PAIMON_DATASKETCHES_PKG_NAME=datasketches-cpp-${PAIMON_DATASKETCHES_BUILD_VERSION}.tar.gz + PAIMON_LUMINA_BUILD_VERSION=0.3.1 PAIMON_LUMINA_BUILD_SHA256_CHECKSUM=2aed1ae238c866c12ee01fad17d52651a9029cba1b972b23cd640959310e2149 PAIMON_LUMINA_PKG_NAME=lumina_release-${PAIMON_LUMINA_BUILD_VERSION}.tar.gz @@ -170,6 +174,7 @@ DEPENDENCIES=( "PAIMON_FMT_URL ${PAIMON_FMT_PKG_NAME} ${THIRDPARTY_MIRROR_URL}https://github.com/fmtlib/fmt/archive/refs/tags/${PAIMON_FMT_BUILD_VERSION}.tar.gz" "PAIMON_GLOG_URL ${PAIMON_GLOG_PKG_NAME} ${THIRDPARTY_MIRROR_URL}https://github.com/google/glog/archive/${PAIMON_GLOG_BUILD_VERSION}.tar.gz" "PAIMON_RAPIDJSON_URL ${PAIMON_RAPIDJSON_PKG_NAME} ${THIRDPARTY_MIRROR_URL}https://github.com/miloyip/rapidjson/archive/${PAIMON_RAPIDJSON_BUILD_VERSION}.tar.gz" + "PAIMON_DATASKETCHES_URL ${PAIMON_DATASKETCHES_PKG_NAME} ${THIRDPARTY_MIRROR_URL}https://github.com/apache/datasketches-cpp/archive/refs/tags/${PAIMON_DATASKETCHES_BUILD_VERSION}.tar.gz" "PAIMON_RE2_URL ${PAIMON_RE2_PKG_NAME} ${THIRDPARTY_MIRROR_URL}https://github.com/google/re2/archive/${PAIMON_RE2_BUILD_VERSION}.tar.gz" "PAIMON_LUMINA_URL ${PAIMON_LUMINA_PKG_NAME} https://paimon-cpp.oss-cn-beijing.aliyuncs.com/thirdparty/lumina/lumina_release-${PAIMON_LUMINA_BUILD_VERSION}.tar.gz" "PAIMON_JINDOSDK_C_LINUX_X86_64_URL ${PAIMON_JINDOSDK_C_LINUX_X86_64_PKG_NAME} https://jindodata-binary.oss-cn-shanghai.aliyuncs.com/release/${PAIMON_JINDOSDK_C_BUILD_VERSION}/jindosdk-${PAIMON_JINDOSDK_C_BUILD_VERSION}-linux.tar.gz"