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 8b93a8b4d..7a9f10395 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 3b97dfd68..d3219ae43 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 @@ -287,6 +289,11 @@ 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_sketch_agg.cpp core/mergetree/compact/aggregate/field_sum_agg.cpp core/mergetree/compact/interval_partition.cpp core/mergetree/compact/loser_tree.cpp @@ -417,6 +424,7 @@ add_paimon_lib(paimon xxhash Threads::Threads RapidJSON + DataSketches ${PAIMON_OBJECT_STORE_LINK_LIBS} STATIC_LINK_LIBS arrow @@ -427,6 +435,7 @@ add_paimon_lib(paimon xxhash Threads::Threads RapidJSON + DataSketches ${PAIMON_OBJECT_STORE_LINK_LIBS} SHARED_LINK_FLAGS ${PAIMON_VERSION_SCRIPT_FLAGS}) @@ -711,14 +720,18 @@ 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_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 @@ -832,6 +845,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 4b51d50c0..e35057eb8 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 aa1e8bcd5..3bdf1e340 100644 --- a/src/paimon/core/core_options.cpp +++ b/src/paimon/core/core_options.cpp @@ -1378,6 +1378,62 @@ 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)); + std::string lower = StringUtils::ToLowerCase(strategy); + if (lower == "merge") { + return NestedKeyNullStrategy::MERGE; + } + if (lower == "ignore") { + return NestedKeyNullStrategy::IGNORE; + } + if (lower == "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 6bc265cf0..c75abbef4 100644 --- a/src/paimon/core/core_options.h +++ b/src/paimon/core/core_options.h @@ -58,6 +58,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, @@ -139,6 +146,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/core_options_test.cpp b/src/paimon/core/core_options_test.cpp index f59649c65..4e43cbcbf 100644 --- a/src/paimon/core/core_options_test.cpp +++ b/src/paimon/core/core_options_test.cpp @@ -17,6 +17,9 @@ #include "paimon/core/core_options.h" #include +#include +#include +#include #include "gtest/gtest.h" #include "paimon/bucket/bucket_function_type.h" @@ -100,6 +103,12 @@ TEST(CoreOptionsTest, TestDefaultValue) { ASSERT_FALSE(core_options.FieldAggIgnoreRetract("f1").value()); ASSERT_EQ(",", core_options.FieldListAggDelimiter("f1").value()); ASSERT_FALSE(core_options.FieldCollectAggDistinct("f1").value()); + ASSERT_TRUE(core_options.FieldNestedUpdateAggNestedKey("f1").value().empty()); + ASSERT_EQ(CoreOptions::NestedKeyNullStrategy::MERGE, + core_options.FieldNestedUpdateAggNestedKeyNullStrategy("f1").value()); + ASSERT_TRUE(core_options.FieldNestedUpdateAggNestedSequenceField("f1").value().empty()); + ASSERT_EQ(std::numeric_limits::max(), + core_options.FieldNestedUpdateAggCountLimit("f1").value()); ASSERT_EQ(MapStorageLayout::DEFAULT, core_options.GetMapStorageLayout("any_col").value()); ASSERT_EQ(256, core_options.GetMapSharedShreddingMaxColumns("any_col").value()); ASSERT_EQ(MapSharedShreddingColumnPlacementPolicy::LRU, @@ -218,6 +227,10 @@ TEST(CoreOptionsTest, TestFromMap) { {"fields.f1.ignore-retract", "true"}, {"fields.f2.list-agg-delimiter", " | "}, {"fields.f2.distinct", "true"}, + {"fields.f3.nested-key", "pk0,pk1"}, + {"fields.f3.nested-key-null-strategy", "ignore"}, + {"fields.f3.nested-sequence-field", "seq0,seq1"}, + {"fields.f3.count-limit", "10"}, {Options::DELETION_VECTORS_ENABLED, "true"}, {Options::DELETION_VECTOR_BITMAP64, "true"}, {Options::DELETION_VECTOR_INDEX_FILE_TARGET_SIZE, "4MB"}, @@ -351,6 +364,13 @@ TEST(CoreOptionsTest, TestFromMap) { ASSERT_TRUE(core_options.FieldAggIgnoreRetract("f1").value()); ASSERT_EQ(" | ", core_options.FieldListAggDelimiter("f2").value()); ASSERT_TRUE(core_options.FieldCollectAggDistinct("f2").value()); + ASSERT_EQ((std::vector{"pk0", "pk1"}), + core_options.FieldNestedUpdateAggNestedKey("f3").value()); + ASSERT_EQ(CoreOptions::NestedKeyNullStrategy::IGNORE, + core_options.FieldNestedUpdateAggNestedKeyNullStrategy("f3").value()); + ASSERT_EQ((std::vector{"seq0", "seq1"}), + core_options.FieldNestedUpdateAggNestedSequenceField("f3").value()); + ASSERT_EQ(10, core_options.FieldNestedUpdateAggCountLimit("f3").value()); ASSERT_TRUE(core_options.DeletionVectorsEnabled()); ASSERT_TRUE(core_options.DeletionVectorsBitmap64()); ASSERT_EQ(4 * 1024 * 1024, core_options.DeletionVectorTargetFileSize()); @@ -477,6 +497,48 @@ TEST(CoreOptionsTest, TestInvalidCase) { ASSERT_NOK_WITH_MSG( CoreOptions::FromMap({{Options::WRITE_SEQUENCE_NUMBER_INIT_MODE, "invalid"}}), "invalid write sequence number init mode: invalid"); + + ASSERT_OK_AND_ASSIGN(CoreOptions invalid_strategy, + CoreOptions::FromMap({{"fields.f0.nested-key-null-strategy", "invalid"}})); + ASSERT_NOK_WITH_MSG(invalid_strategy.FieldNestedUpdateAggNestedKeyNullStrategy("f0"), + "supported values are merge, ignore and error"); + ASSERT_OK_AND_ASSIGN(CoreOptions negative_limit, + CoreOptions::FromMap({{"fields.f0.count-limit", "-1"}})); + ASSERT_NOK_WITH_MSG(negative_limit.FieldNestedUpdateAggCountLimit("f0"), + "must not be negative"); +} + +TEST(CoreOptionsTest, TestNestedKeyNullStrategyIsCaseInsensitive) { + const std::vector> cases = { + {"MERGE", CoreOptions::NestedKeyNullStrategy::MERGE}, + {"Merge", CoreOptions::NestedKeyNullStrategy::MERGE}, + {"IGNORE", CoreOptions::NestedKeyNullStrategy::IGNORE}, + {"Ignore", CoreOptions::NestedKeyNullStrategy::IGNORE}, + {"ERROR", CoreOptions::NestedKeyNullStrategy::ERROR}, + {"eRrOr", CoreOptions::NestedKeyNullStrategy::ERROR}, + }; + for (const auto& [value, expected] : cases) { + ASSERT_OK_AND_ASSIGN(CoreOptions core_options, + CoreOptions::FromMap({{"fields.f0.nested-key-null-strategy", value}})); + ASSERT_EQ(expected, core_options.FieldNestedUpdateAggNestedKeyNullStrategy("f0").value()) + << "value: " << value; + } + + // The rejection message keeps the value exactly as the user wrote it. + ASSERT_OK_AND_ASSIGN(CoreOptions invalid, + CoreOptions::FromMap({{"fields.f0.nested-key-null-strategy", "InVaLid"}})); + ASSERT_NOK_WITH_MSG(invalid.FieldNestedUpdateAggNestedKeyNullStrategy("f0"), + "nested-key-null-strategy: InVaLid"); + + // An absent option keeps the default, but an explicitly empty one is a config error: the + // std::string overload of StringToValue passes "" through, so it reaches the strategy match. + ASSERT_OK_AND_ASSIGN(CoreOptions absent, CoreOptions::FromMap({})); + ASSERT_EQ(CoreOptions::NestedKeyNullStrategy::MERGE, + absent.FieldNestedUpdateAggNestedKeyNullStrategy("f0").value()); + ASSERT_OK_AND_ASSIGN(CoreOptions empty, + CoreOptions::FromMap({{"fields.f0.nested-key-null-strategy", ""}})); + ASSERT_NOK_WITH_MSG(empty.FieldNestedUpdateAggNestedKeyNullStrategy("f0"), + "supported values are merge, ignore and error"); } TEST(CoreOptionsTest, TestLookupCompactMaxIntervalComputedValue) { 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..a42e17d4e 100644 --- a/src/paimon/core/mergetree/compact/aggregate/aggregate_merge_function.cpp +++ b/src/paimon/core/mergetree/compact/aggregate/aggregate_merge_function.cpp @@ -35,7 +35,8 @@ namespace paimon { Result> AggregateMergeFunction::Create( const std::shared_ptr& value_schema, - const std::vector& primary_keys, const CoreOptions& options) { + const std::vector& primary_keys, const CoreOptions& options, + const std::shared_ptr& pool) { std::vector> aggregators; aggregators.reserve(value_schema->num_fields()); for (int32_t i = 0; i < value_schema->num_fields(); i++) { @@ -44,8 +45,8 @@ Result> AggregateMergeFunction::Create( PAIMON_ASSIGN_OR_RAISE(std::string str_agg, GetAggFuncName(field_name, primary_keys, options)); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr agg, - FieldAggregatorFactory::CreateFieldAggregator(field_name, field_type, - str_agg, options)); + FieldAggregatorFactory::CreateFieldAggregator( + field_name, field_type, str_agg, options, pool)); aggregators.push_back(std::move(agg)); } @@ -85,7 +86,7 @@ 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]->Agg(accumulator, input_field)); } row_->SetField(i, merged_field); } diff --git a/src/paimon/core/mergetree/compact/aggregate/aggregate_merge_function.h b/src/paimon/core/mergetree/compact/aggregate/aggregate_merge_function.h index 5e21df4a2..65aa67914 100644 --- a/src/paimon/core/mergetree/compact/aggregate/aggregate_merge_function.h +++ b/src/paimon/core/mergetree/compact/aggregate/aggregate_merge_function.h @@ -47,7 +47,8 @@ class AggregateMergeFunction : public MergeFunction { // value_schema is the schema of parameter value in KeyValue object static Result> Create( const std::shared_ptr& value_schema, - const std::vector& primary_keys, const CoreOptions& options); + const std::vector& primary_keys, const CoreOptions& options, + const std::shared_ptr& pool); void Reset() override { latest_kv_ = std::nullopt; diff --git a/src/paimon/core/mergetree/compact/aggregate/aggregate_merge_function_test.cpp b/src/paimon/core/mergetree/compact/aggregate/aggregate_merge_function_test.cpp index 96b120710..115ecc2db 100644 --- a/src/paimon/core/mergetree/compact/aggregate/aggregate_merge_function_test.cpp +++ b/src/paimon/core/mergetree/compact/aggregate/aggregate_merge_function_test.cpp @@ -84,9 +84,9 @@ TEST(AggregateMergeFunctionTest, TestSimple) { auto value_schema = arrow::schema(fields); ASSERT_OK_AND_ASSIGN(CoreOptions core_options, CoreOptions::FromMap({{Options::FIELDS_DEFAULT_AGG_FUNC, "sum"}})); - ASSERT_OK_AND_ASSIGN( - std::unique_ptr merge_func, - AggregateMergeFunction::Create(value_schema, /*primary_keys=*/{"k0"}, core_options)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr merge_func, + AggregateMergeFunction::Create(value_schema, /*primary_keys=*/{"k0"}, + core_options, GetDefaultPool())); auto pool = GetDefaultPool(); KeyValue kv1(RowKind::Insert(), /*sequence_number=*/0, /*level=*/0, /*key=*/ @@ -124,9 +124,9 @@ TEST(AggregateMergeFunctionTest, TestIgnoreRetract) { ASSERT_OK_AND_ASSIGN(CoreOptions core_options, CoreOptions::FromMap({{Options::FIELDS_DEFAULT_AGG_FUNC, "sum"}, {"fields.v0.ignore-retract", "true"}})); - ASSERT_OK_AND_ASSIGN( - std::unique_ptr merge_func, - AggregateMergeFunction::Create(value_schema, /*primary_keys=*/{"k0"}, core_options)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr merge_func, + AggregateMergeFunction::Create(value_schema, /*primary_keys=*/{"k0"}, + core_options, GetDefaultPool())); auto pool = GetDefaultPool(); KeyValue kv1(RowKind::Insert(), /*sequence_number=*/0, /*level=*/0, /*key=*/ @@ -165,9 +165,9 @@ TEST(AggregateMergeFunctionTest, TestSequenceFields) { ASSERT_OK_AND_ASSIGN(CoreOptions core_options, CoreOptions::FromMap({{Options::SEQUENCE_FIELD, "s0,s1"}, {Options::FIELDS_DEFAULT_AGG_FUNC, "sum"}})); - ASSERT_OK_AND_ASSIGN( - std::unique_ptr merge_func, - AggregateMergeFunction::Create(value_schema, /*primary_keys=*/{"k0"}, core_options)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr merge_func, + AggregateMergeFunction::Create(value_schema, /*primary_keys=*/{"k0"}, + core_options, GetDefaultPool())); auto pool = GetDefaultPool(); // sequence: null, 2 KeyValue kv1( @@ -199,9 +199,9 @@ TEST(AggregateMergeFunctionTest, TestRemoveRecordOnDelete) { CoreOptions core_options, CoreOptions::FromMap({{Options::FIELDS_DEFAULT_AGG_FUNC, "sum"}, {Options::AGGREGATION_REMOVE_RECORD_ON_DELETE, "true"}})); - ASSERT_OK_AND_ASSIGN( - std::unique_ptr merge_func, - AggregateMergeFunction::Create(value_schema, /*primary_keys=*/{"k0"}, core_options)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr merge_func, + AggregateMergeFunction::Create(value_schema, /*primary_keys=*/{"k0"}, + core_options, GetDefaultPool())); auto pool = GetDefaultPool(); @@ -298,9 +298,9 @@ TEST(AggregateMergeFunctionTest, TestDeleteWithoutRemoveRecordOnDelete) { auto value_schema = arrow::schema(fields); ASSERT_OK_AND_ASSIGN(CoreOptions core_options, CoreOptions::FromMap({{Options::FIELDS_DEFAULT_AGG_FUNC, "sum"}})); - ASSERT_OK_AND_ASSIGN( - std::unique_ptr merge_func, - AggregateMergeFunction::Create(value_schema, /*primary_keys=*/{"k0"}, core_options)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr merge_func, + AggregateMergeFunction::Create(value_schema, /*primary_keys=*/{"k0"}, + core_options, GetDefaultPool())); auto pool = GetDefaultPool(); merge_func->Reset(); 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..108a07c2a --- /dev/null +++ b/src/paimon/core/mergetree/compact/aggregate/field_aggregate_utils.cpp @@ -0,0 +1,243 @@ +/* + * 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/types/row_kind.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; + } + PAIMON_ASSIGN_OR_RAISE(const RowKind* lhs_kind, lhs->GetRowKind()); + PAIMON_ASSIGN_OR_RAISE(const RowKind* rhs_kind, rhs->GetRowKind()); + if (lhs_kind != rhs_kind) { + return false; + } + 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, MemoryPool* pool) { + if (DataDefine::GetVariantPtr>(value)) { + return value; + } + std::string_view view = DataDefine::GetStringView(value); + pooled_unique_ptr owned = Bytes::AllocateBytes(view.size(), pool); + 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..ac22079f8 --- /dev/null +++ b/src/paimon/core/mergetree/compact/aggregate/field_aggregate_utils.h @@ -0,0 +1,64 @@ +/* + * 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; +class MemoryPool; + +/// 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. + /// @param pool Pool the copy is allocated from when one is needed. + /// @return The value itself when it already owns its buffer, otherwise an owning copy. + static VariantType OwnedBinary(const VariantType& value, MemoryPool* pool); + + /// 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..013758bec 100644 --- a/src/paimon/core/mergetree/compact/aggregate/field_aggregator.h +++ b/src/paimon/core/mergetree/compact/aggregate/field_aggregator.h @@ -25,21 +25,43 @@ #include "paimon/result.h" namespace paimon { +class MemoryPool; + /// abstract class of aggregating a field of a row. class FieldAggregator { public: virtual ~FieldAggregator() = default; - FieldAggregator(const std::string& name, const std::shared_ptr& field_type) - : name_(name), field_type_(field_type) {} + /// Construct an aggregator for one field. + /// + /// @param name Name of the aggregate function. + /// @param field_type Type of the aggregated field. + /// @param pool Pool every allocation made while aggregating is charged to. Merging runs inside + /// the write and compaction paths, so the caller's pool keeps those bytes visible to memory + /// accounting and to the spill decisions driven by it. + FieldAggregator(const std::string& name, const std::shared_ptr& field_type, + const std::shared_ptr& pool) + : name_(name), field_type_(field_type), pool_(pool) {} - virtual VariantType Agg(const VariantType& accumulator, const VariantType& input_field) = 0; + /// Merge a value into the accumulator. + /// + /// @param accumulator Current aggregated value. + /// @param input_field Value to merge into the accumulator. + /// @return The merged value, or an error Status for aggregators which deserialize external + /// representations and can reject invalid input. + virtual Result Agg(const VariantType& accumulator, + const VariantType& input_field) = 0; /// reset the aggregator to a clean start state. virtual void Reset() {} - virtual VariantType AggReversed(const VariantType& accumulator, - const VariantType& input_field) { + /// Merge an older value into the accumulator. + /// + /// @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 AggReversed(const VariantType& accumulator, + const VariantType& input_field) { return Agg(input_field, accumulator); } @@ -57,9 +79,13 @@ class FieldAggregator { std::shared_ptr GetFieldType() const { return field_type_; } + const std::shared_ptr& GetPool() const { + return pool_; + } protected: std::string name_; std::shared_ptr field_type_; + std::shared_ptr pool_; }; } // namespace paimon 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..9c0f2defb 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,11 @@ #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_sketch_agg.h" #include "paimon/core/mergetree/compact/aggregate/field_sum_agg.h" #include "paimon/result.h" #include "paimon/status.h" @@ -48,33 +52,58 @@ class FieldAggregatorFactory { FieldAggregatorFactory() = delete; ~FieldAggregatorFactory() = delete; + /// Create the aggregator named by @p str_agg for one field. + /// + /// @param field_name Name of the aggregated field. + /// @param field_type Type of the aggregated field. + /// @param str_agg Name of the aggregate function. + /// @param options Table options holding the per-field aggregate settings. + /// @param pool Pool every aggregator allocation is charged to, so that merging done during + /// writes and compaction stays visible to the caller's memory accounting. + /// @return The aggregator, or an error Status for an unknown or misconfigured function. static Result> CreateFieldAggregator( const std::string& field_name, const std::shared_ptr& field_type, - const std::string& str_agg, const CoreOptions& options) { + const std::string& str_agg, const CoreOptions& options, + const std::shared_ptr& pool) { std::unique_ptr field_aggregator; if (str_agg == FieldPrimaryKeyAgg::NAME) { - field_aggregator = std::make_unique(field_type); + field_aggregator = std::make_unique(field_type, pool); } else if (str_agg == FieldLastNonNullValueAgg::NAME) { - field_aggregator = std::make_unique(field_type); + field_aggregator = std::make_unique(field_type, pool); } else if (str_agg == FieldFirstNonNullValueAgg::NAME) { - field_aggregator = std::make_unique(field_type); + field_aggregator = std::make_unique(field_type, pool); } else if (str_agg == FieldLastValueAgg::NAME) { - field_aggregator = std::make_unique(field_type); + field_aggregator = std::make_unique(field_type, pool); } else if (str_agg == FieldFirstValueAgg::NAME) { - field_aggregator = std::make_unique(field_type); + field_aggregator = std::make_unique(field_type, pool); } else if (str_agg == FieldSumAgg::NAME) { - PAIMON_ASSIGN_OR_RAISE(field_aggregator, FieldSumAgg::Create(field_type)); + PAIMON_ASSIGN_OR_RAISE(field_aggregator, FieldSumAgg::Create(field_type, pool)); } else if (str_agg == FieldMinAgg::NAME) { - PAIMON_ASSIGN_OR_RAISE(field_aggregator, FieldMinAgg::Create(field_type)); + PAIMON_ASSIGN_OR_RAISE(field_aggregator, FieldMinAgg::Create(field_type, pool)); } else if (str_agg == FieldMaxAgg::NAME) { - PAIMON_ASSIGN_OR_RAISE(field_aggregator, FieldMaxAgg::Create(field_type)); + PAIMON_ASSIGN_OR_RAISE(field_aggregator, FieldMaxAgg::Create(field_type, pool)); } else if (str_agg == FieldBoolOrAgg::NAME) { - PAIMON_ASSIGN_OR_RAISE(field_aggregator, FieldBoolOrAgg::Create(field_type)); + PAIMON_ASSIGN_OR_RAISE(field_aggregator, FieldBoolOrAgg::Create(field_type, pool)); } else if (str_agg == FieldBoolAndAgg::NAME) { - PAIMON_ASSIGN_OR_RAISE(field_aggregator, FieldBoolAndAgg::Create(field_type)); + PAIMON_ASSIGN_OR_RAISE(field_aggregator, FieldBoolAndAgg::Create(field_type, pool)); } else if (str_agg == FieldListaggAgg::NAME) { PAIMON_ASSIGN_OR_RAISE(field_aggregator, - FieldListaggAgg::Create(field_type, options, field_name)); + FieldListaggAgg::Create(field_type, options, field_name, pool)); + } else if (str_agg == FieldCollectAgg::NAME) { + PAIMON_ASSIGN_OR_RAISE(field_aggregator, + FieldCollectAgg::Create(field_type, options, field_name, pool)); + } else if (str_agg == FieldMergeMapAgg::NAME) { + PAIMON_ASSIGN_OR_RAISE(field_aggregator, + FieldMergeMapAgg::Create(field_type, field_name, pool)); + } else if (str_agg == FieldNestedUpdateAgg::NAME) { + PAIMON_ASSIGN_OR_RAISE(field_aggregator, FieldNestedUpdateAgg::Create( + field_type, options, field_name, pool)); + } else if (str_agg == FieldHllSketchAgg::NAME) { + PAIMON_ASSIGN_OR_RAISE(field_aggregator, + FieldHllSketchAgg::Create(field_type, field_name, pool)); + } else if (str_agg == FieldThetaSketchAgg::NAME) { + PAIMON_ASSIGN_OR_RAISE(field_aggregator, + FieldThetaSketchAgg::Create(field_type, field_name, pool)); } 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..f4351b4ac 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,9 +17,12 @@ #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/memory/memory_pool.h" #include "paimon/testing/utils/testharness.h" namespace paimon::test { @@ -27,80 +30,82 @@ TEST(FieldAggregatorFactoryTest, TestSimple) { { ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({})); ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, - FieldAggregatorFactory::CreateFieldAggregator("f0", arrow::int32(), - "primary-key", options)); + FieldAggregatorFactory::CreateFieldAggregator( + "f0", arrow::int32(), "primary-key", options, GetDefaultPool())); ASSERT_TRUE(dynamic_cast(agg.get())); } { ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({})); - ASSERT_OK_AND_ASSIGN( - std::unique_ptr agg, - FieldAggregatorFactory::CreateFieldAggregator("f0", arrow::int32(), "sum", options)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, + FieldAggregatorFactory::CreateFieldAggregator( + "f0", arrow::int32(), "sum", options, GetDefaultPool())); ASSERT_TRUE(dynamic_cast(agg.get())); } { ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({})); - ASSERT_OK_AND_ASSIGN( - std::unique_ptr agg, - FieldAggregatorFactory::CreateFieldAggregator("f0", arrow::int32(), "min", options)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, + FieldAggregatorFactory::CreateFieldAggregator( + "f0", arrow::int32(), "min", options, GetDefaultPool())); ASSERT_TRUE(dynamic_cast(agg.get())); } { ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({})); - ASSERT_OK_AND_ASSIGN( - std::unique_ptr agg, - FieldAggregatorFactory::CreateFieldAggregator("f0", arrow::int32(), "max", options)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, + FieldAggregatorFactory::CreateFieldAggregator( + "f0", arrow::int32(), "max", options, GetDefaultPool())); ASSERT_TRUE(dynamic_cast(agg.get())); } { ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({})); ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, - FieldAggregatorFactory::CreateFieldAggregator("f0", arrow::boolean(), - "bool_and", options)); + FieldAggregatorFactory::CreateFieldAggregator( + "f0", arrow::boolean(), "bool_and", options, GetDefaultPool())); ASSERT_TRUE(dynamic_cast(agg.get())); } { ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({})); ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, - FieldAggregatorFactory::CreateFieldAggregator("f0", arrow::boolean(), - "bool_or", options)); + FieldAggregatorFactory::CreateFieldAggregator( + "f0", arrow::boolean(), "bool_or", options, GetDefaultPool())); ASSERT_TRUE(dynamic_cast(agg.get())); } { ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({})); - ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, - FieldAggregatorFactory::CreateFieldAggregator( - "f0", arrow::int32(), "last_non_null_value", options)); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr agg, + FieldAggregatorFactory::CreateFieldAggregator( + "f0", arrow::int32(), "last_non_null_value", options, GetDefaultPool())); ASSERT_TRUE(dynamic_cast(agg.get())); } { ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({})); - ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, - FieldAggregatorFactory::CreateFieldAggregator( - "f0", arrow::int32(), "first_non_null_value", options)); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr agg, + FieldAggregatorFactory::CreateFieldAggregator( + "f0", arrow::int32(), "first_non_null_value", options, GetDefaultPool())); ASSERT_TRUE(dynamic_cast(agg.get())); } { ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({})); ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, - FieldAggregatorFactory::CreateFieldAggregator("f0", arrow::int32(), - "last_value", options)); + FieldAggregatorFactory::CreateFieldAggregator( + "f0", arrow::int32(), "last_value", options, GetDefaultPool())); ASSERT_TRUE(dynamic_cast(agg.get())); } { ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({})); ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, - FieldAggregatorFactory::CreateFieldAggregator("f0", arrow::int32(), - "first_value", options)); + FieldAggregatorFactory::CreateFieldAggregator( + "f0", arrow::int32(), "first_value", options, GetDefaultPool())); ASSERT_TRUE(dynamic_cast(agg.get())); } { // test ignore_retract is true ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({{"fields.f0.ignore-retract", "true"}})); - ASSERT_OK_AND_ASSIGN( - std::unique_ptr agg, - FieldAggregatorFactory::CreateFieldAggregator("f0", arrow::int32(), "sum", options)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, + FieldAggregatorFactory::CreateFieldAggregator( + "f0", arrow::int32(), "sum", options, GetDefaultPool())); auto ignore_retract_agg = dynamic_cast(agg.get()); ASSERT_TRUE(ignore_retract_agg); ASSERT_TRUE(dynamic_cast(ignore_retract_agg->agg_.get())); @@ -108,8 +113,8 @@ TEST(FieldAggregatorFactoryTest, TestSimple) { { // test non exist agg ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({})); - auto agg = FieldAggregatorFactory::CreateFieldAggregator("f0", arrow::int32(), - "non-exist-agg", options); + auto agg = FieldAggregatorFactory::CreateFieldAggregator( + "f0", arrow::int32(), "non-exist-agg", options, GetDefaultPool()); ASSERT_FALSE(agg.ok()); } } @@ -119,9 +124,37 @@ TEST(FieldAggregatorFactoryTest, TestRemoveRecordOnDeleteConflictsWithIgnoreRetr CoreOptions options, CoreOptions::FromMap({{Options::AGGREGATION_REMOVE_RECORD_ON_DELETE, "true"}, {"fields.f0.ignore-retract", "true"}})); - ASSERT_NOK_WITH_MSG( - FieldAggregatorFactory::CreateFieldAggregator("f0", arrow::int32(), "sum", options), - "conflicting behavior"); + ASSERT_NOK_WITH_MSG(FieldAggregatorFactory::CreateFieldAggregator("f0", arrow::int32(), "sum", + options, GetDefaultPool()), + "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, GetDefaultPool())); + 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, + GetDefaultPool())); + ASSERT_TRUE(dynamic_cast(merge_map.get())); + ASSERT_OK_AND_ASSIGN(std::unique_ptr nested_update, + FieldAggregatorFactory::CreateFieldAggregator( + "f", nested_type, "nested_update", options, GetDefaultPool())); + ASSERT_TRUE(dynamic_cast(nested_update.get())); + + for (const char* name : {"hll_sketch", "theta_sketch"}) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr aggregator, + FieldAggregatorFactory::CreateFieldAggregator( + "f", arrow::binary(), name, options, GetDefaultPool())); + ASSERT_EQ(name, aggregator->GetName()); + } } } // namespace paimon::test diff --git a/src/paimon/core/mergetree/compact/aggregate/field_bool_agg_test.cpp b/src/paimon/core/mergetree/compact/aggregate/field_bool_agg_test.cpp index a30a1f480..cec2a0e21 100644 --- a/src/paimon/core/mergetree/compact/aggregate/field_bool_agg_test.cpp +++ b/src/paimon/core/mergetree/compact/aggregate/field_bool_agg_test.cpp @@ -21,20 +21,21 @@ #include "paimon/common/data/data_define.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/memory/memory_pool.h" #include "paimon/result.h" #include "paimon/status.h" #include "paimon/testing/utils/testharness.h" namespace paimon::test { TEST(FieldBoolAndAggTest, TestSimple) { - ASSERT_OK_AND_ASSIGN(auto agg, FieldBoolAndAgg::Create(arrow::boolean())); - auto agg_ret = agg->Agg(true, true); + ASSERT_OK_AND_ASSIGN(auto agg, FieldBoolAndAgg::Create(arrow::boolean(), GetDefaultPool())); + auto agg_ret = agg->Agg(true, true).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), true); - agg_ret = agg->Agg(true, false); + agg_ret = agg->Agg(true, false).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), false); - agg_ret = agg->Agg(false, true); + agg_ret = agg->Agg(false, true).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), false); - agg_ret = agg->Agg(false, false); + agg_ret = agg->Agg(false, false).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), false); auto retract_ret = agg->Retract(false, false); @@ -42,35 +43,35 @@ TEST(FieldBoolAndAggTest, TestSimple) { } TEST(FieldBoolAndAggTest, TestInvalidType) { - auto agg = FieldBoolAndAgg::Create(arrow::utf8()); + auto agg = FieldBoolAndAgg::Create(arrow::utf8(), GetDefaultPool()); ASSERT_FALSE(agg.ok()); } TEST(FieldBoolAndAggTest, TestNull) { - ASSERT_OK_AND_ASSIGN(auto agg, FieldBoolAndAgg::Create(arrow::boolean())); + ASSERT_OK_AND_ASSIGN(auto agg, FieldBoolAndAgg::Create(arrow::boolean(), GetDefaultPool())); { - auto agg_ret = agg->Agg(true, NullType()); + auto agg_ret = agg->Agg(true, NullType()).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), true); } { - auto agg_ret = agg->Agg(NullType(), true); + auto agg_ret = agg->Agg(NullType(), true).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), true); } { - auto agg_ret = agg->Agg(NullType(), NullType()); + auto agg_ret = agg->Agg(NullType(), NullType()).value(); ASSERT_TRUE(DataDefine::IsVariantNull(agg_ret)); } } TEST(FieldBoolOrAggTest, TestSimple) { - ASSERT_OK_AND_ASSIGN(auto agg, FieldBoolOrAgg::Create(arrow::boolean())); - auto agg_ret = agg->Agg(true, true); + ASSERT_OK_AND_ASSIGN(auto agg, FieldBoolOrAgg::Create(arrow::boolean(), GetDefaultPool())); + auto agg_ret = agg->Agg(true, true).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), true); - agg_ret = agg->Agg(true, false); + agg_ret = agg->Agg(true, false).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), true); - agg_ret = agg->Agg(false, true); + agg_ret = agg->Agg(false, true).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), true); - agg_ret = agg->Agg(false, false); + agg_ret = agg->Agg(false, false).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), false); auto retract_ret = agg->Retract(false, false); @@ -78,22 +79,22 @@ TEST(FieldBoolOrAggTest, TestSimple) { } TEST(FieldBoolOrAggTest, TestInvalidType) { - auto agg = FieldBoolOrAgg::Create(arrow::utf8()); + auto agg = FieldBoolOrAgg::Create(arrow::utf8(), GetDefaultPool()); ASSERT_FALSE(agg.ok()); } TEST(FieldBoolOrAggTest, TestNull) { - ASSERT_OK_AND_ASSIGN(auto agg, FieldBoolOrAgg::Create(arrow::boolean())); + ASSERT_OK_AND_ASSIGN(auto agg, FieldBoolOrAgg::Create(arrow::boolean(), GetDefaultPool())); { - auto agg_ret = agg->Agg(true, NullType()); + auto agg_ret = agg->Agg(true, NullType()).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), true); } { - auto agg_ret = agg->Agg(NullType(), true); + auto agg_ret = agg->Agg(NullType(), true).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), true); } { - auto agg_ret = agg->Agg(NullType(), NullType()); + auto agg_ret = agg->Agg(NullType(), NullType()).value(); ASSERT_TRUE(DataDefine::IsVariantNull(agg_ret)); } } diff --git a/src/paimon/core/mergetree/compact/aggregate/field_bool_and_agg.h b/src/paimon/core/mergetree/compact/aggregate/field_bool_and_agg.h index 2149d8cc1..c67b317bb 100644 --- a/src/paimon/core/mergetree/compact/aggregate/field_bool_and_agg.h +++ b/src/paimon/core/mergetree/compact/aggregate/field_bool_and_agg.h @@ -26,16 +26,18 @@ namespace paimon { class FieldBoolAndAgg : public FieldAggregator { public: static Result> Create( - const std::shared_ptr& field_type) { + const std::shared_ptr& field_type, + const std::shared_ptr& pool) { if (field_type->id() != arrow::Type::type::BOOL) { return Status::Invalid( fmt::format("invalid field type {} for {}, supposed to be boolean", field_type->ToString(), NAME)); } - return std::unique_ptr(new FieldBoolAndAgg(field_type)); + return std::unique_ptr(new FieldBoolAndAgg(field_type, pool)); } - VariantType Agg(const VariantType& accumulator, const VariantType& input_field) override { + Result Agg(const VariantType& accumulator, + const VariantType& input_field) override { bool accumulator_null = DataDefine::IsVariantNull(accumulator); bool input_null = DataDefine::IsVariantNull(input_field); if (accumulator_null || input_null) { @@ -43,14 +45,15 @@ class FieldBoolAndAgg : public FieldAggregator { } bool accumulator_value = DataDefine::GetVariantValue(accumulator); bool input_value = DataDefine::GetVariantValue(input_field); - return accumulator_value && input_value; + return VariantType(accumulator_value && input_value); } public: static constexpr char NAME[] = "bool_and"; private: - explicit FieldBoolAndAgg(const std::shared_ptr& field_type) - : FieldAggregator(std::string(NAME), field_type) {} + FieldBoolAndAgg(const std::shared_ptr& field_type, + const std::shared_ptr& pool) + : FieldAggregator(std::string(NAME), field_type, pool) {} }; } // namespace paimon diff --git a/src/paimon/core/mergetree/compact/aggregate/field_bool_or_agg.h b/src/paimon/core/mergetree/compact/aggregate/field_bool_or_agg.h index c7a98542c..dab346581 100644 --- a/src/paimon/core/mergetree/compact/aggregate/field_bool_or_agg.h +++ b/src/paimon/core/mergetree/compact/aggregate/field_bool_or_agg.h @@ -26,16 +26,18 @@ namespace paimon { class FieldBoolOrAgg : public FieldAggregator { public: static Result> Create( - const std::shared_ptr& field_type) { + const std::shared_ptr& field_type, + const std::shared_ptr& pool) { if (field_type->id() != arrow::Type::type::BOOL) { return Status::Invalid( fmt::format("invalid field type {} for {}, supposed to be boolean", field_type->ToString(), NAME)); } - return std::unique_ptr(new FieldBoolOrAgg(field_type)); + return std::unique_ptr(new FieldBoolOrAgg(field_type, pool)); } - VariantType Agg(const VariantType& accumulator, const VariantType& input_field) override { + Result Agg(const VariantType& accumulator, + const VariantType& input_field) override { bool accumulator_null = DataDefine::IsVariantNull(accumulator); bool input_null = DataDefine::IsVariantNull(input_field); if (accumulator_null || input_null) { @@ -43,14 +45,15 @@ class FieldBoolOrAgg : public FieldAggregator { } bool accumulator_value = DataDefine::GetVariantValue(accumulator); bool input_value = DataDefine::GetVariantValue(input_field); - return accumulator_value || input_value; + return VariantType(accumulator_value || input_value); } public: static constexpr char NAME[] = "bool_or"; private: - explicit FieldBoolOrAgg(const std::shared_ptr& field_type) - : FieldAggregator(std::string(NAME), field_type) {} + FieldBoolOrAgg(const std::shared_ptr& field_type, + const std::shared_ptr& pool) + : FieldAggregator(std::string(NAME), field_type, pool) {} }; } // namespace paimon 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..163a0cc13 --- /dev/null +++ b/src/paimon/core/mergetree/compact/aggregate/field_collect_agg.cpp @@ -0,0 +1,168 @@ +/* + * 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 { + +// TODO(liangjie.liang): Hash VariantType by its type so that this scan, the key lookup in +// FieldMergeMapAgg and the keyed upsert in FieldNestedUpdateAgg stop being O(n^2). Java only pays +// that cost for constructed element types and uses HashSet/HashMap for the rest. +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, const std::shared_ptr& pool) { + 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, pool)); +} + +Result FieldCollectAgg::Agg(const VariantType& accumulator, + const VariantType& input_field) { + return AggImpl(accumulator, input_field); +} + +Result FieldCollectAgg::AggReversed(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..832186d61 --- /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 + +#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. + /// @param pool Pool the merged arrays are allocated from. + /// @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, const std::shared_ptr& pool); + + Result Agg(const VariantType& accumulator, + const VariantType& input_field) override; + Result AggReversed(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, + const std::shared_ptr& pool) + : FieldAggregator(NAME, field_type, pool), + 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..b9036cad2 --- /dev/null +++ b/src/paimon/core/mergetree/compact/aggregate/field_collect_agg_test.cpp @@ -0,0 +1,293 @@ +/* + * 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 +#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/generic_row.h" +#include "paimon/common/data/serializer/binary_serializer_utils.h" +#include "paimon/core/core_options.h" +#include "paimon/memory/memory_pool.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", GetDefaultPool()); +} + +} // 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->AggReversed(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->Agg(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->Agg(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", GetDefaultPool())); +} + +// Ported from Java FieldAggregatorTest#testFiledCollectAggWith{Row,Array,Map}Type: distinct +// collection over composite element types. +namespace { + +Result> MakeDistinctAgg( + const std::shared_ptr& element_type) { + PAIMON_ASSIGN_OR_RAISE(CoreOptions options, + CoreOptions::FromMap({{"fields.f.distinct", "true"}})); + return FieldCollectAgg::Create(arrow::list(element_type), options, "f", GetDefaultPool()); +} + +VariantType Array(std::vector values) { + return VariantType( + std::static_pointer_cast(std::make_shared(std::move(values)))); +} + +VariantType IntStringRow(int32_t id, std::string_view name) { + std::shared_ptr row = std::make_shared(2); + row->SetField(0, id); + row->SetField(1, name); + return VariantType(std::static_pointer_cast(row)); +} + +VariantType IntStringMap(std::vector> entries) { + std::vector keys; + std::vector values; + for (const auto& entry : entries) { + keys.emplace_back(entry.first); + values.emplace_back(entry.second); + } + return VariantType(std::static_pointer_cast( + std::make_shared(std::make_shared(std::move(keys)), + std::make_shared(std::move(values))))); +} + +/// Decode without going through FieldAggregateUtils, so the assertions stay independent of the +/// equality code under test. +std::vector SortedRows(const VariantType& result) { + auto array = DataDefine::GetVariantValue>(result); + std::vector out; + for (int32_t i = 0; i < array->Size(); ++i) { + std::shared_ptr row = array->GetRow(i, 2); + out.push_back(std::to_string(row->GetInt(0)) + ":" + std::string(row->GetStringView(1))); + } + std::sort(out.begin(), out.end()); + return out; +} + +std::vector SortedArrays(const VariantType& result) { + auto array = DataDefine::GetVariantValue>(result); + std::vector out; + for (int32_t i = 0; i < array->Size(); ++i) { + std::shared_ptr inner = array->GetArray(i); + std::string encoded; + for (int32_t j = 0; j < inner->Size(); ++j) { + encoded += std::to_string(inner->GetInt(j)) + ","; + } + out.push_back(encoded); + } + std::sort(out.begin(), out.end()); + return out; +} + +std::vector SortedMaps(const VariantType& result) { + auto array = DataDefine::GetVariantValue>(result); + std::vector out; + for (int32_t i = 0; i < array->Size(); ++i) { + std::shared_ptr map = array->GetMap(i); + std::shared_ptr keys = map->KeyArray(); + std::shared_ptr values = map->ValueArray(); + std::vector entries; + for (int32_t j = 0; j < map->Size(); ++j) { + entries.push_back(std::to_string(keys->GetInt(j)) + "=" + + std::string(values->GetStringView(j))); + } + std::sort(entries.begin(), entries.end()); + std::string encoded; + for (const std::string& entry : entries) { + encoded += entry + ";"; + } + out.push_back(encoded); + } + std::sort(out.begin(), out.end()); + return out; +} + +} // namespace + +TEST(FieldCollectAggTest, DistinctOverRowElements) { + std::shared_ptr row_type = + arrow::struct_({arrow::field("id", arrow::int32()), arrow::field("name", arrow::utf8())}); + ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, MakeDistinctAgg(row_type)); + + ASSERT_OK_AND_ASSIGN(VariantType empty, + agg->Agg(VariantType(NullType()), VariantType(NullType()))); + ASSERT_TRUE(DataDefine::IsVariantNull(empty)); + + VariantType input1 = Array({IntStringRow(1, "A"), IntStringRow(1, "B")}); + ASSERT_OK_AND_ASSIGN(VariantType first, agg->Agg(VariantType(NullType()), input1)); + ASSERT_EQ((std::vector{"1:A", "1:B"}), SortedRows(first)); + + VariantType input2 = Array({IntStringRow(1, "A"), IntStringRow(2, "A")}); + ASSERT_OK_AND_ASSIGN(VariantType merged, agg->Agg(input1, input2)); + ASSERT_EQ((std::vector{"1:A", "1:B", "2:A"}), SortedRows(merged)); +} + +TEST(FieldCollectAggTest, DistinctOverArrayElements) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, + MakeDistinctAgg(arrow::list(arrow::int32()))); + + ASSERT_OK_AND_ASSIGN(VariantType empty, + agg->Agg(VariantType(NullType()), VariantType(NullType()))); + ASSERT_TRUE(DataDefine::IsVariantNull(empty)); + + VariantType input1 = Array({Array({int32_t{1}, int32_t{1}}), Array({int32_t{1}, int32_t{2}})}); + ASSERT_OK_AND_ASSIGN(VariantType first, agg->Agg(VariantType(NullType()), input1)); + ASSERT_EQ((std::vector{"1,1,", "1,2,"}), SortedArrays(first)); + + VariantType input2 = Array({Array({int32_t{1}, int32_t{1}}), Array({int32_t{1}, int32_t{2}}), + Array({int32_t{2}, int32_t{1}})}); + ASSERT_OK_AND_ASSIGN(VariantType merged, agg->Agg(input1, input2)); + ASSERT_EQ((std::vector{"1,1,", "1,2,", "2,1,"}), SortedArrays(merged)); +} + +TEST(FieldCollectAggTest, DistinctOverMapElements) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, + MakeDistinctAgg(arrow::map(arrow::int32(), arrow::utf8()))); + + ASSERT_OK_AND_ASSIGN(VariantType empty, + agg->Agg(VariantType(NullType()), VariantType(NullType()))); + ASSERT_TRUE(DataDefine::IsVariantNull(empty)); + + VariantType input1 = Array({IntStringMap({{1, "A"}}), IntStringMap({{1, "A"}, {2, "B"}})}); + ASSERT_OK_AND_ASSIGN(VariantType first, agg->Agg(VariantType(NullType()), input1)); + ASSERT_EQ((std::vector{"1=A;", "1=A;2=B;"}), SortedMaps(first)); + + // the second entry has the same content as input1's, only inserted in a different order + VariantType input2 = Array( + {IntStringMap({{1, "A"}}), IntStringMap({{2, "B"}, {1, "A"}}), IntStringMap({{1, "C"}})}); + ASSERT_OK_AND_ASSIGN(VariantType merged, agg->Agg(input1, input2)); + ASSERT_EQ((std::vector{"1=A;", "1=A;2=B;", "1=C;"}), SortedMaps(merged)); +} + +// Ported from Java FieldAggregatorTest#testFieldCollectAggRetractWith{,out}Distinct: retraction +// removes one occurrence per retracted element, for every element type. +TEST(FieldCollectAggTest, RetractRemovesOneOccurrencePerElement) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr int_agg, MakeCollectAgg(true)); + ASSERT_OK_AND_ASSIGN( + VariantType ints, + int_agg->Retract(IntArray({int32_t{1}, int32_t{2}, int32_t{3}}), IntArray({int32_t{1}}))); + ASSERT_EQ((std::vector{2, 3}), Values(ints)); + // duplicates in the accumulator are retracted one at a time + ASSERT_OK_AND_ASSIGN( + VariantType dups, + int_agg->Retract(IntArray({int32_t{1}, int32_t{1}, int32_t{2}, int32_t{2}, int32_t{3}}), + IntArray({int32_t{1}, int32_t{2}, int32_t{3}}))); + ASSERT_EQ((std::vector{1, 2}), Values(dups)); + + std::shared_ptr row_type = + arrow::struct_({arrow::field("id", arrow::int32()), arrow::field("name", arrow::utf8())}); + ASSERT_OK_AND_ASSIGN(std::unique_ptr row_agg, MakeDistinctAgg(row_type)); + ASSERT_OK_AND_ASSIGN(VariantType rows, + row_agg->Retract(Array({IntStringRow(1, "A"), IntStringRow(1, "A"), + IntStringRow(1, "B"), IntStringRow(2, "B")}), + Array({IntStringRow(1, "A"), IntStringRow(2, "B")}))); + ASSERT_EQ((std::vector{"1:A", "1:B"}), SortedRows(rows)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr array_agg, + MakeDistinctAgg(arrow::list(arrow::int32()))); + ASSERT_OK_AND_ASSIGN( + VariantType arrays, + array_agg->Retract( + Array({Array({int32_t{1}, int32_t{1}}), Array({int32_t{1}, int32_t{1}}), + Array({int32_t{1}, int32_t{2}}), Array({int32_t{2}, int32_t{1}})}), + Array({Array({int32_t{1}, int32_t{1}}), Array({int32_t{1}, int32_t{2}})}))); + ASSERT_EQ((std::vector{"1,1,", "2,1,"}), SortedArrays(arrays)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr map_agg, + MakeDistinctAgg(arrow::map(arrow::int32(), arrow::utf8()))); + // the retracted {1=A,2=B} matches the accumulator entry written as {2=B,1=A} + ASSERT_OK_AND_ASSIGN( + VariantType maps, + map_agg->Retract(Array({IntStringMap({{1, "A"}}), IntStringMap({{1, "A"}}), + IntStringMap({{2, "B"}, {1, "A"}}), IntStringMap({{1, "C"}})}), + Array({IntStringMap({{1, "A"}}), IntStringMap({{1, "A"}, {2, "B"}})}))); + ASSERT_EQ((std::vector{"1=A;", "1=C;"}), SortedMaps(maps)); +} + +// Ported from Java FieldAggregatorRetractNullTest: retraction is supported and returns a value. +TEST(FieldCollectAggTest, RetractOnEmptyArraysIsSupported) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, MakeCollectAgg(false)); + ASSERT_OK_AND_ASSIGN(VariantType result, agg->Retract(IntArray({}), IntArray({}))); + ASSERT_TRUE(Values(result).empty()); +} + +} // namespace paimon::test diff --git a/src/paimon/core/mergetree/compact/aggregate/field_first_non_null_value_agg.h b/src/paimon/core/mergetree/compact/aggregate/field_first_non_null_value_agg.h index 611ad2b37..a99077031 100644 --- a/src/paimon/core/mergetree/compact/aggregate/field_first_non_null_value_agg.h +++ b/src/paimon/core/mergetree/compact/aggregate/field_first_non_null_value_agg.h @@ -30,10 +30,12 @@ namespace paimon { /// first non-null value aggregate a field of a row. class FieldFirstNonNullValueAgg : public FieldAggregator { public: - explicit FieldFirstNonNullValueAgg(const std::shared_ptr& field_type) - : FieldAggregator(std::string(NAME), field_type) {} + FieldFirstNonNullValueAgg(const std::shared_ptr& field_type, + const std::shared_ptr& pool) + : FieldAggregator(std::string(NAME), field_type, pool) {} - VariantType Agg(const VariantType& accumulator, const VariantType& input_field) override { + Result Agg(const VariantType& accumulator, + const VariantType& input_field) override { if (!initialized_ && !DataDefine::IsVariantNull(input_field)) { initialized_ = true; return input_field; diff --git a/src/paimon/core/mergetree/compact/aggregate/field_first_non_null_value_agg_test.cpp b/src/paimon/core/mergetree/compact/aggregate/field_first_non_null_value_agg_test.cpp index 0d583b1d6..8cde2ddd4 100644 --- a/src/paimon/core/mergetree/compact/aggregate/field_first_non_null_value_agg_test.cpp +++ b/src/paimon/core/mergetree/compact/aggregate/field_first_non_null_value_agg_test.cpp @@ -20,21 +20,22 @@ #include "arrow/type_fwd.h" #include "gtest/gtest.h" +#include "paimon/memory/memory_pool.h" #include "paimon/result.h" namespace paimon::test { TEST(FieldFirstNonNullValueAggTest, TestSimple) { - auto agg = std::make_unique(arrow::int32()); + auto agg = std::make_unique(arrow::int32(), GetDefaultPool()); - auto agg_ret = agg->Agg(5, 10); + auto agg_ret = agg->Agg(5, 10).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 10); - agg_ret = agg->Agg(10, 20); + agg_ret = agg->Agg(10, 20).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 10); - agg_ret = agg->Agg(10, 30); + agg_ret = agg->Agg(10, 30).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 10); agg->Reset(); - agg_ret = agg->Agg(10, 30); + agg_ret = agg->Agg(10, 30).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 30); auto retract_ret = agg->Retract(10, 30); @@ -42,26 +43,26 @@ TEST(FieldFirstNonNullValueAggTest, TestSimple) { } TEST(FieldFirstNonNullValueAggTest, TestNull) { - auto agg = std::make_unique(arrow::int32()); - auto agg_ret = agg->Agg(5, NullType()); + auto agg = std::make_unique(arrow::int32(), GetDefaultPool()); + auto agg_ret = agg->Agg(5, NullType()).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 5); - agg_ret = agg->Agg(5, 10); + agg_ret = agg->Agg(5, 10).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 10); - agg_ret = agg->Agg(10, NullType()); + agg_ret = agg->Agg(10, NullType()).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 10); - agg_ret = agg->Agg(10, 20); + agg_ret = agg->Agg(10, 20).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 10); agg->Reset(); - agg_ret = agg->Agg(NullType(), NullType()); + agg_ret = agg->Agg(NullType(), NullType()).value(); ASSERT_TRUE(DataDefine::IsVariantNull(agg_ret)); agg->Reset(); - agg_ret = agg->Agg(NullType(), 5); + agg_ret = agg->Agg(NullType(), 5).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 5); - agg_ret = agg->Agg(5, NullType()); + agg_ret = agg->Agg(5, NullType()).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 5); } diff --git a/src/paimon/core/mergetree/compact/aggregate/field_first_value_agg.h b/src/paimon/core/mergetree/compact/aggregate/field_first_value_agg.h index d5fd4dd76..822ae6183 100644 --- a/src/paimon/core/mergetree/compact/aggregate/field_first_value_agg.h +++ b/src/paimon/core/mergetree/compact/aggregate/field_first_value_agg.h @@ -30,10 +30,12 @@ namespace paimon { /// first value aggregate a field of a row. class FieldFirstValueAgg : public FieldAggregator { public: - explicit FieldFirstValueAgg(const std::shared_ptr& field_type) - : FieldAggregator(std::string(NAME), field_type) {} + FieldFirstValueAgg(const std::shared_ptr& field_type, + const std::shared_ptr& pool) + : FieldAggregator(std::string(NAME), field_type, pool) {} - VariantType Agg(const VariantType& accumulator, const VariantType& input_field) override { + Result Agg(const VariantType& accumulator, + const VariantType& input_field) override { if (!initialized_) { initialized_ = true; return input_field; diff --git a/src/paimon/core/mergetree/compact/aggregate/field_first_value_agg_test.cpp b/src/paimon/core/mergetree/compact/aggregate/field_first_value_agg_test.cpp index aab49b021..85cf467be 100644 --- a/src/paimon/core/mergetree/compact/aggregate/field_first_value_agg_test.cpp +++ b/src/paimon/core/mergetree/compact/aggregate/field_first_value_agg_test.cpp @@ -20,21 +20,22 @@ #include "arrow/type_fwd.h" #include "gtest/gtest.h" +#include "paimon/memory/memory_pool.h" #include "paimon/result.h" namespace paimon::test { TEST(FieldFirstValueAggTest, TestSimple) { - auto agg = std::make_unique(arrow::int32()); + auto agg = std::make_unique(arrow::int32(), GetDefaultPool()); - auto agg_ret = agg->Agg(5, 10); + auto agg_ret = agg->Agg(5, 10).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 10); - agg_ret = agg->Agg(10, 20); + agg_ret = agg->Agg(10, 20).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 10); - agg_ret = agg->Agg(10, 30); + agg_ret = agg->Agg(10, 30).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 10); agg->Reset(); - agg_ret = agg->Agg(10, 30); + agg_ret = agg->Agg(10, 30).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 30); auto retract_ret = agg->Retract(10, 30); @@ -42,22 +43,22 @@ TEST(FieldFirstValueAggTest, TestSimple) { } TEST(FieldFirstValueAggTest, TestNull) { - auto agg = std::make_unique(arrow::int32()); - auto agg_ret = agg->Agg(5, NullType()); + auto agg = std::make_unique(arrow::int32(), GetDefaultPool()); + auto agg_ret = agg->Agg(5, NullType()).value(); ASSERT_TRUE(DataDefine::IsVariantNull(agg_ret)); - agg_ret = agg->Agg(NullType(), 10); + agg_ret = agg->Agg(NullType(), 10).value(); ASSERT_TRUE(DataDefine::IsVariantNull(agg_ret)); agg->Reset(); - agg_ret = agg->Agg(NullType(), NullType()); + agg_ret = agg->Agg(NullType(), NullType()).value(); ASSERT_TRUE(DataDefine::IsVariantNull(agg_ret)); agg->Reset(); - agg_ret = agg->Agg(NullType(), 5); + agg_ret = agg->Agg(NullType(), 5).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 5); - agg_ret = agg->Agg(5, NullType()); + agg_ret = agg->Agg(5, NullType()).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 5); } 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..ca73ddd77 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 @@ -28,9 +28,11 @@ namespace paimon { class FieldIgnoreRetractAgg : public FieldAggregator { public: explicit FieldIgnoreRetractAgg(std::unique_ptr&& agg) - : FieldAggregator(agg->GetName(), agg->GetFieldType()), agg_(std::move(agg)) {} + : FieldAggregator(agg->GetName(), agg->GetFieldType(), agg->GetPool()), + agg_(std::move(agg)) {} - VariantType Agg(const VariantType& accumulator, const VariantType& input_field) override { + Result Agg(const VariantType& accumulator, + const VariantType& input_field) override { return agg_->Agg(accumulator, input_field); } 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..dbc089f73 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,37 +17,45 @@ #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/memory/memory_pool.h" #include "paimon/status.h" #include "paimon/testing/utils/testharness.h" namespace paimon::test { TEST(FieldIgnoreRetractAggTest, TestSimple) { - ASSERT_OK_AND_ASSIGN(auto field_sum_agg, FieldSumAgg::Create(arrow::int32())); + ASSERT_OK_AND_ASSIGN(auto field_sum_agg, FieldSumAgg::Create(arrow::int32(), GetDefaultPool())); auto agg = std::make_unique(std::move(field_sum_agg)); - auto agg_ret = agg->Agg(5, 10); + auto agg_ret = agg->Agg(5, 10).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 15); ASSERT_OK_AND_ASSIGN(auto retract_ret, agg->Retract(5, 10)); ASSERT_EQ(DataDefine::GetVariantValue(retract_ret), 5); } TEST(FieldIgnoreRetractAggTest, TestNull) { - ASSERT_OK_AND_ASSIGN(auto field_sum_agg, FieldSumAgg::Create(arrow::int32())); + ASSERT_OK_AND_ASSIGN(auto field_sum_agg, FieldSumAgg::Create(arrow::int32(), GetDefaultPool())); auto agg = std::make_unique(std::move(field_sum_agg)); { - auto agg_ret = agg->Agg(5, NullType()); + auto agg_ret = agg->Agg(5, NullType()).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 5); } { - auto agg_ret = agg->Agg(NullType(), 10); + auto agg_ret = agg->Agg(NullType(), 10).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 10); } { - auto agg_ret = agg->Agg(NullType(), NullType()); + auto agg_ret = agg->Agg(NullType(), NullType()).value(); ASSERT_TRUE(DataDefine::IsVariantNull(agg_ret)); } { @@ -64,4 +72,27 @@ 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", GetDefaultPool())); + 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->AggReversed(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_last_non_null_value_agg.h b/src/paimon/core/mergetree/compact/aggregate/field_last_non_null_value_agg.h index 29edc4e11..b0ccc876a 100644 --- a/src/paimon/core/mergetree/compact/aggregate/field_last_non_null_value_agg.h +++ b/src/paimon/core/mergetree/compact/aggregate/field_last_non_null_value_agg.h @@ -31,10 +31,12 @@ namespace paimon { /// last non-null value aggregate a field of a row. class FieldLastNonNullValueAgg : public FieldAggregator { public: - explicit FieldLastNonNullValueAgg(const std::shared_ptr& field_type) - : FieldAggregator(std::string(NAME), field_type) {} + FieldLastNonNullValueAgg(const std::shared_ptr& field_type, + const std::shared_ptr& pool) + : FieldAggregator(std::string(NAME), field_type, pool) {} - VariantType Agg(const VariantType& accumulator, const VariantType& input_field) override { + Result Agg(const VariantType& accumulator, + const VariantType& input_field) override { return DataDefine::IsVariantNull(input_field) ? accumulator : input_field; } diff --git a/src/paimon/core/mergetree/compact/aggregate/field_last_non_null_value_agg_test.cpp b/src/paimon/core/mergetree/compact/aggregate/field_last_non_null_value_agg_test.cpp index 47564ca6c..e60f93610 100644 --- a/src/paimon/core/mergetree/compact/aggregate/field_last_non_null_value_agg_test.cpp +++ b/src/paimon/core/mergetree/compact/aggregate/field_last_non_null_value_agg_test.cpp @@ -20,31 +20,32 @@ #include "arrow/type_fwd.h" #include "gtest/gtest.h" +#include "paimon/memory/memory_pool.h" #include "paimon/status.h" #include "paimon/testing/utils/testharness.h" namespace paimon::test { TEST(FieldLastNonNullValueAggTest, TestSimple) { - auto agg = std::make_unique(arrow::int32()); + auto agg = std::make_unique(arrow::int32(), GetDefaultPool()); - auto agg_ret = agg->Agg(5, 10); + auto agg_ret = agg->Agg(5, 10).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 10); ASSERT_OK_AND_ASSIGN(auto retract_ret, agg->Retract(5, 10)); ASSERT_TRUE(DataDefine::IsVariantNull(retract_ret)); } TEST(FieldLastNonNullValueAggTest, TestNull) { - auto agg = std::make_unique(arrow::int32()); + auto agg = std::make_unique(arrow::int32(), GetDefaultPool()); { - auto agg_ret = agg->Agg(5, NullType()); + auto agg_ret = agg->Agg(5, NullType()).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 5); } { - auto agg_ret = agg->Agg(NullType(), 10); + auto agg_ret = agg->Agg(NullType(), 10).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 10); } { - auto agg_ret = agg->Agg(NullType(), NullType()); + auto agg_ret = agg->Agg(NullType(), NullType()).value(); ASSERT_TRUE(DataDefine::IsVariantNull(agg_ret)); } diff --git a/src/paimon/core/mergetree/compact/aggregate/field_last_value_agg.h b/src/paimon/core/mergetree/compact/aggregate/field_last_value_agg.h index 0001eacb6..a813b1ff8 100644 --- a/src/paimon/core/mergetree/compact/aggregate/field_last_value_agg.h +++ b/src/paimon/core/mergetree/compact/aggregate/field_last_value_agg.h @@ -32,10 +32,12 @@ namespace paimon { /// last value aggregate a field of a row. class FieldLastValueAgg : public FieldAggregator { public: - explicit FieldLastValueAgg(const std::shared_ptr& field_type) - : FieldAggregator(std::string(NAME), field_type) {} + FieldLastValueAgg(const std::shared_ptr& field_type, + const std::shared_ptr& pool) + : FieldAggregator(std::string(NAME), field_type, pool) {} - VariantType Agg(const VariantType& accumulator, const VariantType& input_field) override { + Result Agg(const VariantType& accumulator, + const VariantType& input_field) override { return input_field; } diff --git a/src/paimon/core/mergetree/compact/aggregate/field_last_value_agg_test.cpp b/src/paimon/core/mergetree/compact/aggregate/field_last_value_agg_test.cpp index 832d0091a..765466ef2 100644 --- a/src/paimon/core/mergetree/compact/aggregate/field_last_value_agg_test.cpp +++ b/src/paimon/core/mergetree/compact/aggregate/field_last_value_agg_test.cpp @@ -20,14 +20,15 @@ #include "arrow/type_fwd.h" #include "gtest/gtest.h" +#include "paimon/memory/memory_pool.h" #include "paimon/status.h" #include "paimon/testing/utils/testharness.h" namespace paimon::test { TEST(FieldLastValueAggTest, TestSimple) { - auto agg = std::make_unique(arrow::int32()); + auto agg = std::make_unique(arrow::int32(), GetDefaultPool()); - auto agg_ret = agg->Agg(5, 10); + auto agg_ret = agg->Agg(5, 10).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 10); ASSERT_OK_AND_ASSIGN(auto retract_ret, agg->Retract(5, 10)); @@ -35,17 +36,17 @@ TEST(FieldLastValueAggTest, TestSimple) { } TEST(FieldLastValueAggTest, TestNull) { - auto agg = std::make_unique(arrow::int32()); + auto agg = std::make_unique(arrow::int32(), GetDefaultPool()); { - auto agg_ret = agg->Agg(5, NullType()); + auto agg_ret = agg->Agg(5, NullType()).value(); ASSERT_TRUE(DataDefine::IsVariantNull(agg_ret)); } { - auto agg_ret = agg->Agg(NullType(), 10); + auto agg_ret = agg->Agg(NullType(), 10).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 10); } { - auto agg_ret = agg->Agg(NullType(), NullType()); + auto agg_ret = agg->Agg(NullType(), NullType()).value(); ASSERT_TRUE(DataDefine::IsVariantNull(agg_ret)); } diff --git a/src/paimon/core/mergetree/compact/aggregate/field_listagg_agg.h b/src/paimon/core/mergetree/compact/aggregate/field_listagg_agg.h index ff6ac149d..e381c0625 100644 --- a/src/paimon/core/mergetree/compact/aggregate/field_listagg_agg.h +++ b/src/paimon/core/mergetree/compact/aggregate/field_listagg_agg.h @@ -19,6 +19,7 @@ #include #include #include +#include #include "paimon/common/data/data_define.h" #include "paimon/core/core_options.h" @@ -37,7 +38,7 @@ class FieldListaggAgg : public FieldAggregator { static Result> Create( const std::shared_ptr& field_type, const CoreOptions& options, - const std::string& field_name) { + const std::string& field_name, const std::shared_ptr& pool) { if (field_type->id() != arrow::Type::type::STRING) { return Status::Invalid( fmt::format("invalid field type {} for field '{}' of {}, supposed to be string", @@ -50,10 +51,11 @@ class FieldListaggAgg : public FieldAggregator { delimiter = " "; } return std::unique_ptr( - new FieldListaggAgg(field_type, std::move(delimiter), distinct)); + new FieldListaggAgg(field_type, std::move(delimiter), distinct, pool)); } - VariantType Agg(const VariantType& accumulator, const VariantType& input_field) override { + Result Agg(const VariantType& accumulator, + const VariantType& input_field) override { bool accumulator_null = DataDefine::IsVariantNull(accumulator); bool input_null = DataDefine::IsVariantNull(input_field); if (accumulator_null || input_null) { @@ -79,7 +81,7 @@ class FieldListaggAgg : public FieldAggregator { new_result.append(in_str); result_ = std::move(new_result); } - return std::string_view{result_}; + return VariantType(std::string_view{result_}); } private: @@ -121,9 +123,9 @@ class FieldListaggAgg : public FieldAggregator { return result; } - explicit FieldListaggAgg(const std::shared_ptr& field_type, - std::string delimiter, bool distinct) - : FieldAggregator(std::string(NAME), field_type), + FieldListaggAgg(const std::shared_ptr& field_type, std::string delimiter, + bool distinct, const std::shared_ptr& pool) + : FieldAggregator(std::string(NAME), field_type, pool), delimiter_(std::move(delimiter)), distinct_(distinct) {} diff --git a/src/paimon/core/mergetree/compact/aggregate/field_listagg_agg_test.cpp b/src/paimon/core/mergetree/compact/aggregate/field_listagg_agg_test.cpp index 6c1a6ffd3..30de32a69 100644 --- a/src/paimon/core/mergetree/compact/aggregate/field_listagg_agg_test.cpp +++ b/src/paimon/core/mergetree/compact/aggregate/field_listagg_agg_test.cpp @@ -23,6 +23,7 @@ #include "arrow/type_fwd.h" #include "gtest/gtest.h" #include "paimon/core/core_options.h" +#include "paimon/memory/memory_pool.h" #include "paimon/status.h" #include "paimon/testing/utils/testharness.h" @@ -36,19 +37,19 @@ class FieldListaggAggTest : public testing::Test { opts["fields.f.list-agg-delimiter"] = delimiter; opts["fields.f.distinct"] = distinct ? "true" : "false"; PAIMON_ASSIGN_OR_RAISE(auto options, CoreOptions::FromMap(opts)); - return FieldListaggAgg::Create(arrow::utf8(), std::move(options), "f"); + return FieldListaggAgg::Create(arrow::utf8(), std::move(options), "f", GetDefaultPool()); } }; TEST_F(FieldListaggAggTest, TestSimple) { ASSERT_OK_AND_ASSIGN(auto agg, MakeAgg()); - auto ret = agg->Agg(std::string_view("hello"), std::string_view(" world")); + auto ret = agg->Agg(std::string_view("hello"), std::string_view(" world")).value(); ASSERT_EQ(DataDefine::GetVariantValue(ret), "hello, world"); } TEST_F(FieldListaggAggTest, TestDelimiter) { ASSERT_OK_AND_ASSIGN(auto agg, MakeAgg("-")); - auto ret = agg->Agg(std::string_view("user1"), std::string_view("user2")); + auto ret = agg->Agg(std::string_view("user1"), std::string_view("user2")).value(); ASSERT_EQ(DataDefine::GetVariantValue(ret), "user1-user2"); } @@ -57,17 +58,17 @@ TEST_F(FieldListaggAggTest, TestNull) { // input null -> return accumulator { - auto ret = agg->Agg(std::string_view("hello"), NullType()); + auto ret = agg->Agg(std::string_view("hello"), NullType()).value(); ASSERT_EQ(DataDefine::GetVariantValue(ret), "hello"); } // accumulator null -> return input { - auto ret = agg->Agg(NullType(), std::string_view("world")); + auto ret = agg->Agg(NullType(), std::string_view("world")).value(); ASSERT_EQ(DataDefine::GetVariantValue(ret), "world"); } // both null -> return null { - auto ret = agg->Agg(NullType(), NullType()); + auto ret = agg->Agg(NullType(), NullType()).value(); ASSERT_TRUE(DataDefine::IsVariantNull(ret)); } } @@ -77,17 +78,17 @@ TEST_F(FieldListaggAggTest, TestEmptyString) { // empty input -> return accumulator { - auto ret = agg->Agg(std::string_view("hello"), std::string_view("")); + auto ret = agg->Agg(std::string_view("hello"), std::string_view("")).value(); ASSERT_EQ(DataDefine::GetVariantValue(ret), "hello"); } // empty accumulator -> return input { - auto ret = agg->Agg(std::string_view(""), std::string_view("world")); + auto ret = agg->Agg(std::string_view(""), std::string_view("world")).value(); ASSERT_EQ(DataDefine::GetVariantValue(ret), "world"); } // both empty -> return input (which is empty) { - auto ret = agg->Agg(std::string_view(""), std::string_view("")); + auto ret = agg->Agg(std::string_view(""), std::string_view("")).value(); ASSERT_EQ(DataDefine::GetVariantValue(ret), ""); } } @@ -96,9 +97,9 @@ TEST_F(FieldListaggAggTest, TestMultipleAccumulation) { ASSERT_OK_AND_ASSIGN(auto agg, MakeAgg()); // "a" + "," + "b" = "a,b", then "a,b" + "," + "c" = "a,b,c" - auto ret = agg->Agg(std::string_view("a"), std::string_view("b")); + auto ret = agg->Agg(std::string_view("a"), std::string_view("b")).value(); ASSERT_EQ(DataDefine::GetVariantValue(ret), "a,b"); - ret = agg->Agg(std::move(ret), std::string_view("c")); + ret = agg->Agg(std::move(ret), std::string_view("c")).value(); ASSERT_EQ(DataDefine::GetVariantValue(ret), "a,b,c"); } @@ -106,7 +107,7 @@ TEST_F(FieldListaggAggTest, TestDistinct) { ASSERT_OK_AND_ASSIGN(auto agg, MakeAgg(";", true)); // "a;b" + "b;c" -> "a;b;c" (deduplicate "b") - auto ret = agg->Agg(std::string_view("a;b"), std::string_view("b;c")); + auto ret = agg->Agg(std::string_view("a;b"), std::string_view("b;c")).value(); ASSERT_EQ(DataDefine::GetVariantValue(ret), "a;b;c"); } @@ -114,7 +115,7 @@ TEST_F(FieldListaggAggTest, TestDistinctNoDuplicates) { ASSERT_OK_AND_ASSIGN(auto agg, MakeAgg(" ", true)); // "a b" + "c d" -> "a b c d" (no dups to remove) - auto ret = agg->Agg(std::string_view("a b"), std::string_view("c d")); + auto ret = agg->Agg(std::string_view("a b"), std::string_view("c d")).value(); ASSERT_EQ(DataDefine::GetVariantValue(ret), "a b c d"); } @@ -122,7 +123,7 @@ TEST_F(FieldListaggAggTest, TestDistinctWithEmptyDelimiterFallsBackToWhitespace) ASSERT_OK_AND_ASSIGN(auto agg, MakeAgg("", true)); // Empty delimiter falls back to whitespace, so the repeated "b" is removed. - auto ret = agg->Agg(std::string_view("a b"), std::string_view("b c")); + auto ret = agg->Agg(std::string_view("a b"), std::string_view("b c")).value(); ASSERT_EQ(DataDefine::GetVariantValue(ret), "a b c"); } @@ -130,7 +131,7 @@ TEST_F(FieldListaggAggTest, TestDistinctEmptyInput) { ASSERT_OK_AND_ASSIGN(auto agg, MakeAgg(";", true)); // empty input -> return accumulator - auto ret = agg->Agg(std::string_view("a;b"), std::string_view("")); + auto ret = agg->Agg(std::string_view("a;b"), std::string_view("")).value(); ASSERT_EQ(DataDefine::GetVariantValue(ret), "a;b"); } @@ -138,13 +139,13 @@ TEST_F(FieldListaggAggTest, TestDistinctFalse) { ASSERT_OK_AND_ASSIGN(auto agg, MakeAgg(";", false)); // "a;b" + "b;c" -> "a;b;b;c" (no dedup) - auto ret = agg->Agg(std::string_view("a;b"), std::string_view("b;c")); + auto ret = agg->Agg(std::string_view("a;b"), std::string_view("b;c")).value(); ASSERT_EQ(DataDefine::GetVariantValue(ret), "a;b;b;c"); } TEST_F(FieldListaggAggTest, TestInvalidType) { EXPECT_OK_AND_ASSIGN(auto options, CoreOptions::FromMap({})); - auto result = FieldListaggAgg::Create(arrow::int32(), options, "f"); + auto result = FieldListaggAgg::Create(arrow::int32(), options, "f", GetDefaultPool()); ASSERT_FALSE(result.ok()); ASSERT_TRUE(result.status().ToString().find("supposed to be string") != std::string::npos) << result.status().ToString(); diff --git a/src/paimon/core/mergetree/compact/aggregate/field_max_agg.h b/src/paimon/core/mergetree/compact/aggregate/field_max_agg.h index 6435b4af1..447f53902 100644 --- a/src/paimon/core/mergetree/compact/aggregate/field_max_agg.h +++ b/src/paimon/core/mergetree/compact/aggregate/field_max_agg.h @@ -27,12 +27,14 @@ namespace paimon { class FieldMaxAgg : public FieldAggregator { public: static Result> Create( - const std::shared_ptr& field_type) { + const std::shared_ptr& field_type, + const std::shared_ptr& pool) { PAIMON_ASSIGN_OR_RAISE(FieldMaxFunc max_func, CreateMaxFunc(field_type)); - return std::unique_ptr(new FieldMaxAgg(field_type, max_func)); + return std::unique_ptr(new FieldMaxAgg(field_type, max_func, pool)); } - VariantType Agg(const VariantType& accumulator, const VariantType& input_field) override { + Result Agg(const VariantType& accumulator, + const VariantType& input_field) override { bool accumulator_null = DataDefine::IsVariantNull(accumulator); bool input_null = DataDefine::IsVariantNull(input_field); if (accumulator_null || input_null) { @@ -48,8 +50,9 @@ class FieldMaxAgg : public FieldAggregator { using FieldMaxFunc = std::function; - FieldMaxAgg(const std::shared_ptr& field_type, const FieldMaxFunc& max_func) - : FieldAggregator(std::string(NAME), field_type), max_func_(max_func) {} + FieldMaxAgg(const std::shared_ptr& field_type, const FieldMaxFunc& max_func, + const std::shared_ptr& pool) + : FieldAggregator(std::string(NAME), field_type, pool), max_func_(max_func) {} static Result CreateMaxFunc(const std::shared_ptr& field_type) { arrow::Type::type type = field_type->id(); 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..188f71e9f --- /dev/null +++ b/src/paimon/core/mergetree/compact/aggregate/field_merge_map_agg.cpp @@ -0,0 +1,153 @@ +/* + * 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, + const std::shared_ptr& pool) { + 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(), pool)); +} + +Result FieldMergeMapAgg::Agg(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..6d2e884b6 --- /dev/null +++ b/src/paimon/core/mergetree/compact/aggregate/field_merge_map_agg.h @@ -0,0 +1,63 @@ +/* + * 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/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. + /// @param pool Pool the merged maps are allocated from. + /// @return A merge_map aggregator, or an error Status. + static Result> Create( + const std::shared_ptr& field_type, const std::string& field_name, + const std::shared_ptr& pool); + + Result Agg(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, + const std::shared_ptr& pool) + : FieldAggregator(NAME, field_type, pool), + 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..998f5c1d5 --- /dev/null +++ b/src/paimon/core/mergetree/compact/aggregate/field_merge_map_agg_test.cpp @@ -0,0 +1,98 @@ +/* + * 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/memory/memory_pool.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", + GetDefaultPool())); + + ASSERT_OK_AND_ASSIGN(VariantType merged, + agg->Agg(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", + GetDefaultPool())); + VariantType map = IntMap({int32_t{1}}, {int32_t{10}}); + ASSERT_OK_AND_ASSIGN(VariantType result, agg->Agg(VariantType(NullType()), map)); + ASSERT_EQ(10, FindValue(result, 1)); + ASSERT_NOK(FieldMergeMapAgg::Create(arrow::int32(), "f", GetDefaultPool())); +} + +// Ported from Java FieldAggregatorRetractNullTest: retraction is supported and returns a value. +TEST(FieldMergeMapAggTest, RetractOnEmptyMapIsSupported) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, + FieldMergeMapAgg::Create(arrow::map(arrow::int32(), arrow::int32()), "f", + GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(VariantType result, agg->Retract(IntMap({}, {}), IntMap({}, {}))); + ASSERT_EQ(0, DataDefine::GetVariantValue>(result)->Size()); +} + +} // namespace paimon::test diff --git a/src/paimon/core/mergetree/compact/aggregate/field_min_agg.h b/src/paimon/core/mergetree/compact/aggregate/field_min_agg.h index b456c0f73..b973cd617 100644 --- a/src/paimon/core/mergetree/compact/aggregate/field_min_agg.h +++ b/src/paimon/core/mergetree/compact/aggregate/field_min_agg.h @@ -27,12 +27,14 @@ namespace paimon { class FieldMinAgg : public FieldAggregator { public: static Result> Create( - const std::shared_ptr& field_type) { + const std::shared_ptr& field_type, + const std::shared_ptr& pool) { PAIMON_ASSIGN_OR_RAISE(FieldMinFunc min_func, CreateMinFunc(field_type)); - return std::unique_ptr(new FieldMinAgg(field_type, min_func)); + return std::unique_ptr(new FieldMinAgg(field_type, min_func, pool)); } - VariantType Agg(const VariantType& accumulator, const VariantType& input_field) override { + Result Agg(const VariantType& accumulator, + const VariantType& input_field) override { bool accumulator_null = DataDefine::IsVariantNull(accumulator); bool input_null = DataDefine::IsVariantNull(input_field); if (accumulator_null || input_null) { @@ -48,8 +50,9 @@ class FieldMinAgg : public FieldAggregator { using FieldMinFunc = std::function; - FieldMinAgg(const std::shared_ptr& field_type, const FieldMinFunc& min_func) - : FieldAggregator(std::string(NAME), field_type), min_func_(min_func) {} + FieldMinAgg(const std::shared_ptr& field_type, const FieldMinFunc& min_func, + const std::shared_ptr& pool) + : FieldAggregator(std::string(NAME), field_type, pool), min_func_(min_func) {} static Result CreateMinFunc(const std::shared_ptr& field_type) { arrow::Type::type type = field_type->id(); diff --git a/src/paimon/core/mergetree/compact/aggregate/field_min_max_agg_test.cpp b/src/paimon/core/mergetree/compact/aggregate/field_min_max_agg_test.cpp index f54c75e4a..6d48aaa97 100644 --- a/src/paimon/core/mergetree/compact/aggregate/field_min_max_agg_test.cpp +++ b/src/paimon/core/mergetree/compact/aggregate/field_min_max_agg_test.cpp @@ -32,6 +32,7 @@ #include "paimon/core/mergetree/compact/aggregate/field_min_agg.h" #include "paimon/data/decimal.h" #include "paimon/data/timestamp.h" +#include "paimon/memory/memory_pool.h" #include "paimon/result.h" #include "paimon/status.h" #include "paimon/testing/utils/testharness.h" @@ -62,15 +63,15 @@ void CheckJavaCompatibleFloatingPointMinMax(const std::shared_ptr values = {-infinity, -static_cast(0.0), static_cast(0.0), infinity, nan}; - ASSERT_OK_AND_ASSIGN(auto field_min_agg, FieldMinAgg::Create(type)); - ASSERT_OK_AND_ASSIGN(auto field_max_agg, FieldMaxAgg::Create(type)); + ASSERT_OK_AND_ASSIGN(auto field_min_agg, FieldMinAgg::Create(type, GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto field_max_agg, FieldMaxAgg::Create(type, GetDefaultPool())); for (size_t i = 0; i < values.size(); ++i) { for (size_t j = 0; j < values.size(); ++j) { - VariantType min_result = field_min_agg->Agg(values[i], values[j]); + ASSERT_OK_AND_ASSIGN(VariantType min_result, field_min_agg->Agg(values[i], values[j])); AssertSameFloatingPoint(DataDefine::GetVariantValue(min_result), values[std::min(i, j)]); - VariantType max_result = field_max_agg->Agg(values[i], values[j]); + ASSERT_OK_AND_ASSIGN(VariantType max_result, field_max_agg->Agg(values[i], values[j])); AssertSameFloatingPoint(DataDefine::GetVariantValue(max_result), values[std::max(i, j)]); } @@ -79,13 +80,15 @@ void CheckJavaCompatibleFloatingPointMinMax(const std::shared_ptrAgg(5, 10); + ASSERT_OK_AND_ASSIGN(auto field_min_agg, + FieldMinAgg::Create(arrow::int32(), GetDefaultPool())); + auto agg_ret = field_min_agg->Agg(5, 10).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 5); } { - ASSERT_OK_AND_ASSIGN(auto field_max_agg, FieldMaxAgg::Create(arrow::int32())); - auto agg_ret = field_max_agg->Agg(5, 10); + ASSERT_OK_AND_ASSIGN(auto field_max_agg, + FieldMaxAgg::Create(arrow::int32(), GetDefaultPool())); + auto agg_ret = field_max_agg->Agg(5, 10).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 10); } } @@ -96,31 +99,31 @@ TEST(FieldMinMaxAggTest, TestJavaCompatibleFloatingPointOrder) { } TEST(FieldMinMaxAggTest, TestInvalidType) { - auto field_min_agg = FieldMinAgg::Create(arrow::boolean()); + auto field_min_agg = FieldMinAgg::Create(arrow::boolean(), GetDefaultPool()); ASSERT_FALSE(field_min_agg.ok()); - auto field_max_agg = FieldMaxAgg::Create(arrow::boolean()); + auto field_max_agg = FieldMaxAgg::Create(arrow::boolean(), GetDefaultPool()); ASSERT_FALSE(field_max_agg.ok()); } TEST(FieldMinMaxAggTest, TestNull) { - ASSERT_OK_AND_ASSIGN(auto field_min_agg, FieldMinAgg::Create(arrow::int32())); - ASSERT_OK_AND_ASSIGN(auto field_max_agg, FieldMaxAgg::Create(arrow::int32())); + ASSERT_OK_AND_ASSIGN(auto field_min_agg, FieldMinAgg::Create(arrow::int32(), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto field_max_agg, FieldMaxAgg::Create(arrow::int32(), GetDefaultPool())); { - auto agg_ret = field_min_agg->Agg(5, NullType()); + auto agg_ret = field_min_agg->Agg(5, NullType()).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 5); - agg_ret = field_max_agg->Agg(5, NullType()); + agg_ret = field_max_agg->Agg(5, NullType()).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 5); } { - auto agg_ret = field_min_agg->Agg(NullType(), 10); + auto agg_ret = field_min_agg->Agg(NullType(), 10).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 10); - agg_ret = field_max_agg->Agg(NullType(), 10); + agg_ret = field_max_agg->Agg(NullType(), 10).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 10); } { - auto agg_ret = field_min_agg->Agg(NullType(), NullType()); + auto agg_ret = field_min_agg->Agg(NullType(), NullType()).value(); ASSERT_TRUE(DataDefine::IsVariantNull(agg_ret)); - agg_ret = field_max_agg->Agg(NullType(), NullType()); + agg_ret = field_max_agg->Agg(NullType(), NullType()).value(); ASSERT_TRUE(DataDefine::IsVariantNull(agg_ret)); } } @@ -128,16 +131,16 @@ TEST(FieldMinMaxAggTest, TestNull) { TEST(FieldMinMaxAggTest, TestVariantType) { auto CheckResult = [](const std::shared_ptr& type, const VariantType& large, const VariantType& small) { - ASSERT_OK_AND_ASSIGN(auto field_min_agg, FieldMinAgg::Create(type)); - ASSERT_OK_AND_ASSIGN(auto field_max_agg, FieldMaxAgg::Create(type)); - auto agg_ret = field_min_agg->Agg(small, large); + ASSERT_OK_AND_ASSIGN(auto field_min_agg, FieldMinAgg::Create(type, GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(auto field_max_agg, FieldMaxAgg::Create(type, GetDefaultPool())); + auto agg_ret = field_min_agg->Agg(small, large).value(); ASSERT_EQ(agg_ret, small); - agg_ret = field_min_agg->Agg(large, small); + agg_ret = field_min_agg->Agg(large, small).value(); ASSERT_EQ(agg_ret, small); - agg_ret = field_max_agg->Agg(small, large); + agg_ret = field_max_agg->Agg(small, large).value(); ASSERT_EQ(agg_ret, large); - agg_ret = field_max_agg->Agg(large, small); + agg_ret = field_max_agg->Agg(large, small).value(); ASSERT_EQ(agg_ret, large); }; 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..5b5e04422 --- /dev/null +++ b/src/paimon/core/mergetree/compact/aggregate/field_nested_update_agg.cpp @@ -0,0 +1,345 @@ +/* + * 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 "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/types/row_kind.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())); + } + 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, + const std::shared_ptr& pool) + : FieldAggregator(NAME, field_type, pool), + 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, const std::shared_ptr& pool) { + 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, pool)); +} + +Result FieldNestedUpdateAgg::Agg(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::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, FieldAggregateUtils::Equals(VariantType(*iter), + VariantType(retract_row), row_type_)); + 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..7093c2b30 --- /dev/null +++ b/src/paimon/core/mergetree/compact/aggregate/field_nested_update_agg.h @@ -0,0 +1,77 @@ +/* + * 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. + /// @param pool Pool the merged nested rows are allocated from. + /// @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, const std::shared_ptr& pool); + + Result Agg(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, + const std::shared_ptr& pool); + + 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; + + 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..069753f34 --- /dev/null +++ b/src/paimon/core/mergetree/compact/aggregate/field_nested_update_agg_test.cpp @@ -0,0 +1,487 @@ +/* + * 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 + +#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/common/types/row_kind.h" +#include "paimon/core/core_options.h" +#include "paimon/memory/memory_pool.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", GetDefaultPool()); +} + +} // 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->Agg(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->Agg(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->Agg(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->Agg(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)); +} + +// matches Java's RecordEqualiser, which compares the row kind before any field +TEST(FieldNestedUpdateAggTest, RetractRequiresMatchingRowKind) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, MakeAgg({})); + std::shared_ptr retract_row = std::make_shared(3); + retract_row->SetField(0, int32_t{1}); + retract_row->SetField(1, 1); + retract_row->SetField(2, 10); + retract_row->SetRowKind(RowKind::Delete()); + + ASSERT_OK_AND_ASSIGN( + VariantType kept, + agg->Retract(Rows({Row(int32_t{1}, 1, 10)}), + Rows({VariantType(std::static_pointer_cast(retract_row))}))); + ASSERT_EQ(1, GetRows(kept)->Size()); + ASSERT_TRUE(FindRow(kept, 1)); +} + +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->Agg(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->Agg(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", GetDefaultPool())); + + 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", GetDefaultPool())); + + 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", GetDefaultPool())); + + 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", GetDefaultPool())); + + ASSERT_OK_AND_ASSIGN(CoreOptions negative_limit, + CoreOptions::FromMap({{"fields.f.count-limit", "-1"}})); + ASSERT_NOK(FieldNestedUpdateAgg::Create(NestedType(), negative_limit, "f", GetDefaultPool())); + + // Java resolves nested-key names with List.indexOf and accepts repeats, so we must too + ASSERT_OK_AND_ASSIGN(CoreOptions repeated_key, + CoreOptions::FromMap({{"fields.f.nested-key", "id,id"}})); + ASSERT_OK(FieldNestedUpdateAgg::Create(NestedType(), repeated_key, "f", GetDefaultPool())); +} + +// Ported from Java FieldAggregatorTest: composite nested keys, multiple sequence fields and the +// count-limit / null-key-strategy boundaries. +namespace { + +std::shared_ptr CompositeKeyType() { + return arrow::list( + arrow::struct_({arrow::field("k0", arrow::int32()), arrow::field("k1", arrow::int32()), + arrow::field("v", arrow::utf8()), arrow::field("seq", arrow::int32()), + arrow::field("seq2", arrow::int32())})); +} + +VariantType KeyedRow(VariantType k0, VariantType k1, std::string_view v, int32_t seq, + int32_t seq2) { + std::shared_ptr row = std::make_shared(5); + row->SetField(0, k0); + row->SetField(1, k1); + row->SetField(2, v); + row->SetField(3, seq); + row->SetField(4, seq2); + return VariantType(std::static_pointer_cast(row)); +} + +Result> MakeKeyedAgg( + const std::map& options_map) { + PAIMON_ASSIGN_OR_RAISE(CoreOptions options, CoreOptions::FromMap(options_map)); + return FieldNestedUpdateAgg::Create(CompositeKeyType(), options, "f", GetDefaultPool()); +} + +std::vector SortedKeyed(const VariantType& value) { + std::shared_ptr rows = GetRows(value); + std::vector out; + for (int32_t i = 0; i < rows->Size(); ++i) { + std::shared_ptr row = rows->GetRow(i, 5); + std::string k0 = row->IsNullAt(0) ? "null" : std::to_string(row->GetInt(0)); + std::string k1 = row->IsNullAt(1) ? "null" : std::to_string(row->GetInt(1)); + out.push_back(k0 + "/" + k1 + "/" + std::string(row->GetStringView(2))); + } + std::sort(out.begin(), out.end()); + return out; +} + +} // namespace + +TEST(FieldNestedUpdateAggTest, CountLimitStillUpdatesExistingCompositeKey) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, + MakeKeyedAgg({{"fields.f.nested-key", "k0,k1"}, + {"fields.f.nested-sequence-field", "seq"}, + {"fields.f.count-limit", "2"}})); + VariantType acc = VariantType(NullType()); + ASSERT_OK_AND_ASSIGN(acc, agg->Agg(acc, Rows({KeyedRow(0, 1, "B", 1, 0)}))); + ASSERT_OK_AND_ASSIGN(acc, agg->Agg(acc, Rows({KeyedRow(1, 2, "C", 3, 0)}))); + + // at the limit an existing key can still be updated + ASSERT_OK_AND_ASSIGN(acc, agg->Agg(acc, Rows({KeyedRow(0, 1, "B_updated", 4, 0)}))); + ASSERT_EQ((std::vector{"0/1/B_updated", "1/2/C"}), SortedKeyed(acc)); + + // but a new key is rejected + ASSERT_OK_AND_ASSIGN(acc, agg->Agg(acc, Rows({KeyedRow(2, 3, "D", 5, 0)}))); + ASSERT_EQ((std::vector{"0/1/B_updated", "1/2/C"}), SortedKeyed(acc)); +} + +TEST(FieldNestedUpdateAggTest, MultipleSequenceFieldsCompareInOrder) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, + MakeKeyedAgg({{"fields.f.nested-key", "k0,k1"}, + {"fields.f.nested-sequence-field", "seq,seq2"}})); + VariantType acc = VariantType(NullType()); + ASSERT_OK_AND_ASSIGN(acc, agg->Agg(acc, Rows({KeyedRow(0, 1, "A", 1, 5)}))); + + // same leading sequence, smaller second field, so the row is kept + ASSERT_OK_AND_ASSIGN(acc, agg->Agg(acc, Rows({KeyedRow(0, 1, "older", 1, 4)}))); + ASSERT_EQ((std::vector{"0/1/A"}), SortedKeyed(acc)); + + // same leading sequence, larger second field, so the row wins + ASSERT_OK_AND_ASSIGN(acc, agg->Agg(acc, Rows({KeyedRow(0, 1, "newer", 1, 6)}))); + ASSERT_EQ((std::vector{"0/1/newer"}), SortedKeyed(acc)); +} + +TEST(FieldNestedUpdateAggTest, NullKeyStrategyAppliesToRetractInput) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr merge_agg, + MakeKeyedAgg({{"fields.f.nested-key", "k0,k1"}})); + VariantType acc = VariantType(NullType()); + ASSERT_OK_AND_ASSIGN(acc, merge_agg->Agg(acc, Rows({KeyedRow(0, 0, "A", 0, 0)}))); + ASSERT_OK_AND_ASSIGN(acc, merge_agg->Agg(acc, Rows({KeyedRow(1, 1, "B", 0, 0)}))); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr ignore_agg, + MakeKeyedAgg({{"fields.f.nested-key", "k0,k1"}, + {"fields.f.nested-key-null-strategy", "ignore"}})); + ASSERT_OK_AND_ASSIGN( + VariantType kept, + ignore_agg->Retract(acc, Rows({KeyedRow(0, VariantType(NullType()), "X", 0, 0)}))); + ASSERT_EQ((std::vector{"0/0/A", "1/1/B"}), SortedKeyed(kept)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr error_agg, + MakeKeyedAgg({{"fields.f.nested-key", "k0,k1"}, + {"fields.f.nested-key-null-strategy", "error"}})); + ASSERT_NOK(error_agg->Retract(acc, Rows({KeyedRow(0, VariantType(NullType()), "X", 0, 0)}))); +} + +// Ported from Java FieldAggregatorTest#testFieldNestedAppendAgg*: without a nested key rows are +// appended rather than upserted, and retraction removes an equal row. +TEST(FieldNestedUpdateAggTest, AppendsWithoutNestedKey) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, MakeKeyedAgg({})); + VariantType acc = VariantType(NullType()); + ASSERT_OK_AND_ASSIGN(acc, agg->Agg(acc, Rows({KeyedRow(0, 1, "B", 0, 0)}))); + ASSERT_EQ((std::vector{"0/1/B"}), SortedKeyed(acc)); + + // same key fields but a different value, so it is appended instead of replacing + ASSERT_OK_AND_ASSIGN(acc, agg->Agg(acc, Rows({KeyedRow(0, 1, "b", 0, 0)}))); + ASSERT_EQ((std::vector{"0/1/B", "0/1/b"}), SortedKeyed(acc)); + + ASSERT_OK_AND_ASSIGN(acc, agg->Retract(acc, Rows({KeyedRow(0, 1, "b", 0, 0)}))); + ASSERT_EQ((std::vector{"0/1/B"}), SortedKeyed(acc)); +} + +TEST(FieldNestedUpdateAggTest, AppendsWithoutNestedKeyRespectCountLimit) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, + MakeKeyedAgg({{"fields.f.count-limit", "2"}})); + VariantType acc = VariantType(NullType()); + ASSERT_OK_AND_ASSIGN(acc, agg->Agg(acc, Rows({KeyedRow(0, 1, "B", 0, 0)}))); + ASSERT_OK_AND_ASSIGN(acc, agg->Agg(acc, Rows({KeyedRow(0, 1, "b", 0, 0)}))); + ASSERT_EQ((std::vector{"0/1/B", "0/1/b"}), SortedKeyed(acc)); + + ASSERT_OK_AND_ASSIGN(acc, agg->Agg(acc, Rows({KeyedRow(0, 1, "C", 0, 0)}))); + ASSERT_EQ((std::vector{"0/1/B", "0/1/b"}), SortedKeyed(acc)); + + // the limit also applies within a single input array, and null elements are skipped + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_input_agg, + MakeKeyedAgg({{"fields.f.count-limit", "2"}})); + ASSERT_OK_AND_ASSIGN( + VariantType first, + first_input_agg->Agg(VariantType(NullType()), + Rows({KeyedRow(0, 1, "B", 0, 0), VariantType(NullType()), + KeyedRow(0, 1, "b", 0, 0), KeyedRow(0, 1, "C", 0, 0)}))); + ASSERT_EQ((std::vector{"0/1/B", "0/1/b"}), SortedKeyed(first)); +} + +TEST(FieldNestedUpdateAggTest, CountLimitAppliesWithinFirstInputArray) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr with_seq, + MakeKeyedAgg({{"fields.f.nested-key", "k0,k1"}, + {"fields.f.nested-sequence-field", "seq"}, + {"fields.f.count-limit", "2"}})); + ASSERT_OK_AND_ASSIGN( + VariantType seq_result, + with_seq->Agg(VariantType(NullType()), + Rows({KeyedRow(0, 1, "B", 1, 0), KeyedRow(1, 2, "C", 3, 0), + KeyedRow(2, 3, "D", 5, 0), KeyedRow(0, 1, "B_updated", 4, 0)}))); + ASSERT_EQ((std::vector{"0/1/B_updated", "1/2/C"}), SortedKeyed(seq_result)); + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr without_seq, + MakeKeyedAgg({{"fields.f.nested-key", "k0,k1"}, {"fields.f.count-limit", "2"}})); + ASSERT_OK_AND_ASSIGN( + VariantType no_seq_result, + without_seq->Agg(VariantType(NullType()), + Rows({KeyedRow(0, 1, "B", 0, 0), KeyedRow(1, 2, "C", 0, 0), + KeyedRow(2, 3, "D", 0, 0), KeyedRow(0, 1, "B_updated", 0, 0)}))); + ASSERT_EQ((std::vector{"0/1/B_updated", "1/2/C"}), SortedKeyed(no_seq_result)); +} + +TEST(FieldNestedUpdateAggTest, CountLimitStillUpdatesExistingKeyWithoutSequence) { + ASSERT_OK_AND_ASSIGN( + std::unique_ptr agg, + MakeKeyedAgg({{"fields.f.nested-key", "k0,k1"}, {"fields.f.count-limit", "2"}})); + VariantType acc = VariantType(NullType()); + ASSERT_OK_AND_ASSIGN(acc, agg->Agg(acc, Rows({KeyedRow(0, 1, "B", 0, 0)}))); + ASSERT_OK_AND_ASSIGN(acc, agg->Agg(acc, Rows({KeyedRow(1, 2, "C", 0, 0)}))); + ASSERT_OK_AND_ASSIGN(acc, agg->Agg(acc, Rows({KeyedRow(0, 1, "B_updated", 0, 0)}))); + ASSERT_EQ((std::vector{"0/1/B_updated", "1/2/C"}), SortedKeyed(acc)); + + ASSERT_OK_AND_ASSIGN(acc, agg->Agg(acc, Rows({KeyedRow(2, 3, "D", 0, 0)}))); + ASSERT_EQ((std::vector{"0/1/B_updated", "1/2/C"}), SortedKeyed(acc)); +} + +// MERGE keeps rows whose nested key is partially or fully null, treating null as a key value. +TEST(FieldNestedUpdateAggTest, MergeStrategyKeepsNullNestedKeys) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, + MakeKeyedAgg({{"fields.f.nested-key", "k0,k1"}})); + VariantType null_k1 = VariantType(NullType()); + + ASSERT_OK_AND_ASSIGN(VariantType partial, agg->Agg(VariantType(NullType()), + Rows({KeyedRow(0, null_k1, "C", 3, 0)}))); + ASSERT_EQ((std::vector{"0/null/C"}), SortedKeyed(partial)); + + ASSERT_OK_AND_ASSIGN(VariantType full, agg->Agg(VariantType(NullType()), + Rows({KeyedRow(null_k1, null_k1, "D", 4, 0)}))); + ASSERT_EQ((std::vector{"null/null/D"}), SortedKeyed(full)); + + VariantType acc = VariantType(NullType()); + ASSERT_OK_AND_ASSIGN(acc, agg->Agg(acc, Rows({KeyedRow(0, 0, "A", 1, 0)}))); + ASSERT_OK_AND_ASSIGN(acc, agg->Agg(acc, Rows({KeyedRow(0, 1, "B", 2, 0)}))); + ASSERT_OK_AND_ASSIGN(acc, agg->Agg(acc, Rows({KeyedRow(0, null_k1, "C", 3, 0)}))); + ASSERT_EQ((std::vector{"0/0/A", "0/1/B", "0/null/C"}), SortedKeyed(acc)); +} + +// Null-keyed rows consume the count limit under MERGE but are skipped entirely under IGNORE. +TEST(FieldNestedUpdateAggTest, CountLimitInteractsWithNullKeyStrategies) { + VariantType null_key = VariantType(NullType()); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr merge_agg, + MakeKeyedAgg({{"fields.f.nested-key", "k0,k1"}, + {"fields.f.nested-sequence-field", "seq"}, + {"fields.f.count-limit", "3"}})); + VariantType merged = VariantType(NullType()); + ASSERT_OK_AND_ASSIGN(merged, merge_agg->Agg(merged, Rows({KeyedRow(0, 1, "B", 1, 0)}))); + ASSERT_OK_AND_ASSIGN(merged, + merge_agg->Agg(merged, Rows({KeyedRow(null_key, 2, "NULL_2", 2, 0)}))); + ASSERT_OK_AND_ASSIGN( + merged, merge_agg->Agg(merged, Rows({KeyedRow(null_key, null_key, "NULL_NULL", 3, 0)}))); + ASSERT_OK_AND_ASSIGN(merged, merge_agg->Agg(merged, Rows({KeyedRow(1, 2, "C", 5, 0)}))); + ASSERT_OK_AND_ASSIGN(merged, merge_agg->Agg(merged, Rows({KeyedRow(0, 1, "B_updated", 4, 0)}))); + ASSERT_EQ((std::vector{"0/1/B_updated", "null/2/NULL_2", "null/null/NULL_NULL"}), + SortedKeyed(merged)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr ignore_agg, + MakeKeyedAgg({{"fields.f.nested-key", "k0,k1"}, + {"fields.f.nested-key-null-strategy", "ignore"}, + {"fields.f.nested-sequence-field", "seq"}, + {"fields.f.count-limit", "3"}})); + VariantType ignored = VariantType(NullType()); + ASSERT_OK_AND_ASSIGN(ignored, ignore_agg->Agg(ignored, Rows({KeyedRow(0, 1, "B", 1, 0)}))); + ASSERT_OK_AND_ASSIGN(ignored, + ignore_agg->Agg(ignored, Rows({KeyedRow(null_key, 2, "NULL_2", 2, 0)}))); + ASSERT_OK_AND_ASSIGN( + ignored, ignore_agg->Agg(ignored, Rows({KeyedRow(null_key, null_key, "NN", 3, 0)}))); + ASSERT_OK_AND_ASSIGN(ignored, ignore_agg->Agg(ignored, Rows({KeyedRow(1, 2, "C", 3, 0)}))); + ASSERT_OK_AND_ASSIGN(ignored, + ignore_agg->Agg(ignored, Rows({KeyedRow(0, 1, "B_updated", 4, 0)}))); + ASSERT_EQ((std::vector{"0/1/B_updated", "1/2/C"}), SortedKeyed(ignored)); + + // room is left for a third real key + ASSERT_OK_AND_ASSIGN(ignored, ignore_agg->Agg(ignored, Rows({KeyedRow(2, 3, "D", 5, 0)}))); + ASSERT_EQ((std::vector{"0/1/B_updated", "1/2/C", "2/3/D"}), SortedKeyed(ignored)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr error_agg, + MakeKeyedAgg({{"fields.f.nested-key", "k0,k1"}, + {"fields.f.nested-key-null-strategy", "error"}, + {"fields.f.count-limit", "3"}})); + ASSERT_NOK( + error_agg->Agg(VariantType(NullType()), Rows({KeyedRow(null_key, 2, "NULL_2", 2, 0)}))); +} + +TEST(FieldNestedUpdateAggTest, NullKeyStrategyAppliesToRetractAccumulator) { + VariantType null_key = VariantType(NullType()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr merge_agg, + MakeKeyedAgg({{"fields.f.nested-key", "k0,k1"}})); + VariantType acc = VariantType(NullType()); + ASSERT_OK_AND_ASSIGN(acc, merge_agg->Agg(acc, Rows({KeyedRow(0, 0, "A", 0, 0)}))); + ASSERT_OK_AND_ASSIGN(acc, merge_agg->Agg(acc, Rows({KeyedRow(null_key, 1, "N", 0, 0)}))); + + // IGNORE drops the null-keyed accumulator row while retracting an unrelated key + ASSERT_OK_AND_ASSIGN(std::unique_ptr ignore_agg, + MakeKeyedAgg({{"fields.f.nested-key", "k0,k1"}, + {"fields.f.nested-key-null-strategy", "ignore"}})); + ASSERT_OK_AND_ASSIGN(VariantType result, + ignore_agg->Retract(acc, Rows({KeyedRow(9, 9, "X", 0, 0)}))); + ASSERT_EQ((std::vector{"0/0/A"}), SortedKeyed(result)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr error_agg, + MakeKeyedAgg({{"fields.f.nested-key", "k0,k1"}, + {"fields.f.nested-key-null-strategy", "error"}})); + ASSERT_NOK(error_agg->Retract(acc, Rows({KeyedRow(9, 9, "X", 0, 0)}))); +} + +// Ported from Java FieldAggregatorRetractNullTest: retraction is supported and returns a value. +TEST(FieldNestedUpdateAggTest, RetractOnEmptyArraysIsSupported) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, MakeKeyedAgg({})); + ASSERT_OK_AND_ASSIGN(VariantType result, agg->Retract(Rows({}), Rows({}))); + ASSERT_EQ(0, GetRows(result)->Size()); +} + +} // namespace paimon::test diff --git a/src/paimon/core/mergetree/compact/aggregate/field_primary_key_agg.h b/src/paimon/core/mergetree/compact/aggregate/field_primary_key_agg.h index 6f6a72679..f60288523 100644 --- a/src/paimon/core/mergetree/compact/aggregate/field_primary_key_agg.h +++ b/src/paimon/core/mergetree/compact/aggregate/field_primary_key_agg.h @@ -31,10 +31,12 @@ namespace paimon { /// primary key aggregate a field of a row. class FieldPrimaryKeyAgg : public FieldAggregator { public: - explicit FieldPrimaryKeyAgg(const std::shared_ptr& field_type) - : FieldAggregator(std::string(NAME), field_type) {} + FieldPrimaryKeyAgg(const std::shared_ptr& field_type, + const std::shared_ptr& pool) + : FieldAggregator(std::string(NAME), field_type, pool) {} - VariantType Agg(const VariantType& accumulator, const VariantType& input_field) override { + Result Agg(const VariantType& accumulator, + const VariantType& input_field) override { return input_field; } diff --git a/src/paimon/core/mergetree/compact/aggregate/field_primary_key_agg_test.cpp b/src/paimon/core/mergetree/compact/aggregate/field_primary_key_agg_test.cpp index 1a8350702..27c345f78 100644 --- a/src/paimon/core/mergetree/compact/aggregate/field_primary_key_agg_test.cpp +++ b/src/paimon/core/mergetree/compact/aggregate/field_primary_key_agg_test.cpp @@ -20,32 +20,33 @@ #include "arrow/type_fwd.h" #include "gtest/gtest.h" +#include "paimon/memory/memory_pool.h" #include "paimon/status.h" #include "paimon/testing/utils/testharness.h" namespace paimon::test { // just for test, in practice, for primary key, accumulator will always equals to input_field TEST(FieldPrimaryKeyAggTest, TestSimple) { - auto agg = std::make_unique(arrow::int32()); + auto agg = std::make_unique(arrow::int32(), GetDefaultPool()); - auto agg_ret = agg->Agg(5, 10); + auto agg_ret = agg->Agg(5, 10).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 10); ASSERT_OK_AND_ASSIGN(auto retract_ret, agg->Retract(5, 10)); ASSERT_EQ(DataDefine::GetVariantValue(retract_ret), 10); } TEST(FieldPrimaryKeyAggTest, TestNull) { - auto agg = std::make_unique(arrow::int32()); + auto agg = std::make_unique(arrow::int32(), GetDefaultPool()); { - auto agg_ret = agg->Agg(5, NullType()); + auto agg_ret = agg->Agg(5, NullType()).value(); ASSERT_TRUE(DataDefine::IsVariantNull(agg_ret)); } { - auto agg_ret = agg->Agg(NullType(), 10); + auto agg_ret = agg->Agg(NullType(), 10).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 10); } { - auto agg_ret = agg->Agg(NullType(), NullType()); + auto agg_ret = agg->Agg(NullType(), NullType()).value(); ASSERT_TRUE(DataDefine::IsVariantNull(agg_ret)); } 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..caf334023 --- /dev/null +++ b/src/paimon/core/mergetree/compact/aggregate/field_sketch_agg.cpp @@ -0,0 +1,137 @@ +/* + * 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, MemoryPool* pool) { + pooled_unique_ptr result = Bytes::AllocateBytes(serialized.size() * sizeof(T), pool); + 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, + const std::shared_ptr& pool) { + PAIMON_RETURN_NOT_OK(ValidateSketchType(field_type, field_name, NAME)); + return std::unique_ptr(new FieldHllSketchAgg(field_type, pool)); +} + +Result FieldHllSketchAgg::Agg(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 VariantType(NullType()); + } + if (accumulator_null || input_null) { + // AggReversed swaps the arguments, so either side may be the row-owned accumulator + return FieldAggregateUtils::OwnedBinary(accumulator_null ? input_field : accumulator, + pool_.get()); + } + 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(), pool_.get())); + } 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, + const std::shared_ptr& pool) { + PAIMON_RETURN_NOT_OK(ValidateSketchType(field_type, field_name, NAME)); + return std::unique_ptr(new FieldThetaSketchAgg(field_type, pool)); +} + +Result FieldThetaSketchAgg::Agg(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 VariantType(NullType()); + } + if (accumulator_null || input_null) { + // AggReversed swaps the arguments, so either side may be the row-owned accumulator + return FieldAggregateUtils::OwnedBinary(accumulator_null ? input_field : accumulator, + pool_.get()); + } + 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(), pool_.get())); + } 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..9754b6061 --- /dev/null +++ b/src/paimon/core/mergetree/compact/aggregate/field_sketch_agg.h @@ -0,0 +1,74 @@ +/* + * 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. + /// @param pool Pool the unioned sketch bytes are allocated from. + /// @return An hll_sketch aggregator, or an error Status. + static Result> Create( + const std::shared_ptr& field_type, const std::string& field_name, + const std::shared_ptr& pool); + + Result Agg(const VariantType& accumulator, + const VariantType& input_field) override; + + private: + FieldHllSketchAgg(const std::shared_ptr& field_type, + const std::shared_ptr& pool) + : FieldAggregator(NAME, field_type, pool) {} +}; + +/// 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. + /// @param pool Pool the unioned sketch bytes are allocated from. + /// @return A theta_sketch aggregator, or an error Status. + static Result> Create( + const std::shared_ptr& field_type, const std::string& field_name, + const std::shared_ptr& pool); + + Result Agg(const VariantType& accumulator, + const VariantType& input_field) override; + + private: + FieldThetaSketchAgg(const std::shared_ptr& field_type, + const std::shared_ptr& pool) + : FieldAggregator(NAME, field_type, pool) {} +}; + +} // 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..ab1d1abb4 --- /dev/null +++ b/src/paimon/core/mergetree/compact/aggregate/field_sketch_agg_test.cpp @@ -0,0 +1,210 @@ +/* + * 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" + +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)); +} + +// 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("hll", arrow::binary()), + arrow::field("theta", arrow::binary())}; + ASSERT_OK_AND_ASSIGN( + CoreOptions options, + CoreOptions::FromMap({{"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, GetDefaultPool())); + + MemoryPool* pool = GetDefaultPool().get(); + ASSERT_OK( + merge_func->Add(KeyValue(RowKind::Insert(), /*sequence_number=*/0, /*level=*/0, + BinaryRowGenerator::GenerateRowPtr({10}, pool), + BinaryRowGenerator::GenerateRowPtr( + {10, 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, 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()}, pool)))); + + ASSERT_OK_AND_ASSIGN(std::optional result, merge_func->GetResult()); + ASSERT_TRUE(result.has_value()); + std::vector> scribbled = ScribbleFreedMemory(); + + std::string_view hll_bytes = result->value->GetStringView(1); + 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(2); + 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", GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(VariantType result, agg->Agg(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", GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(VariantType result, agg->Agg(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()); +} + +// AggReversed swaps its arguments, so either side may carry the row-owned accumulator +TEST(FieldSketchAggTest, NullArgumentsKeepOwnershipInBothDirections) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr agg, + FieldHllSketchAgg::Create(arrow::binary(), "f", GetDefaultPool())); + VariantType null_value = VariantType(NullType()); + // pass a view, not an owning value, or OwnedBinary short-circuits and skips the copy path + VariantType owned = Hll({1, 2, 3}); + VariantType sketch = VariantType(DataDefine::GetStringView(owned)); + + for (const VariantType& result : + {agg->Agg(sketch, null_value).value(), agg->Agg(null_value, sketch).value(), + agg->AggReversed(sketch, null_value).value(), + agg->AggReversed(null_value, sketch).value()}) { + ASSERT_TRUE(DataDefine::GetVariantPtr>(result)) + << "the surviving value must own its buffer"; + std::string_view bytes = DataDefine::GetStringView(result); + ASSERT_NEAR( + 3.0, datasketches::hll_sketch::deserialize(bytes.data(), bytes.size()).get_estimate(), + 0.1); + } + + ASSERT_OK_AND_ASSIGN(VariantType both_null, agg->Agg(null_value, null_value)); + ASSERT_TRUE(DataDefine::IsVariantNull(both_null)); + ASSERT_OK_AND_ASSIGN(VariantType both_null_reversed, agg->AggReversed(null_value, null_value)); + ASSERT_TRUE(DataDefine::IsVariantNull(both_null_reversed)); +} + +// Java leaves retract unimplemented for the sketch aggregators, so it must surface as an error +// pointing at fields..ignore-retract rather than silently succeeding. +TEST(FieldSketchAggTest, RetractionIsUnsupported) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr hll_agg, + FieldHllSketchAgg::Create(arrow::binary(), "f", GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(std::unique_ptr theta_agg, + FieldThetaSketchAgg::Create(arrow::binary(), "f", GetDefaultPool())); + ASSERT_NOK_WITH_MSG(hll_agg->Retract(Hll({1}), Hll({1})), "does not support retraction"); + ASSERT_NOK_WITH_MSG(theta_agg->Retract(Theta({1}), Theta({1})), "does not support retraction"); +} + +TEST(FieldSketchAggTest, ReportsInvalidBytesAndTypes) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr hll_agg, + FieldHllSketchAgg::Create(arrow::binary(), "f", GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(std::unique_ptr theta_agg, + FieldThetaSketchAgg::Create(arrow::binary(), "f", GetDefaultPool())); + VariantType invalid = VariantType(std::string_view("bad")); + ASSERT_NOK(hll_agg->Agg(invalid, Hll({1}))); + ASSERT_NOK(hll_agg->AggReversed(invalid, Hll({1}))); + ASSERT_NOK(theta_agg->Agg(invalid, Theta({1}))); + ASSERT_NOK(FieldHllSketchAgg::Create(arrow::int32(), "f", GetDefaultPool())); + ASSERT_NOK(FieldThetaSketchAgg::Create(arrow::int32(), "f", GetDefaultPool())); +} + +} // namespace paimon::test diff --git a/src/paimon/core/mergetree/compact/aggregate/field_sum_agg.h b/src/paimon/core/mergetree/compact/aggregate/field_sum_agg.h index c93ce310d..f4173848d 100644 --- a/src/paimon/core/mergetree/compact/aggregate/field_sum_agg.h +++ b/src/paimon/core/mergetree/compact/aggregate/field_sum_agg.h @@ -35,13 +35,15 @@ namespace paimon { class FieldSumAgg : public FieldAggregator { public: static Result> Create( - const std::shared_ptr& field_type) { + const std::shared_ptr& field_type, + const std::shared_ptr& pool) { PAIMON_ASSIGN_OR_RAISE(FieldSumFunc sum_func, CreateSumFunc(field_type)); PAIMON_ASSIGN_OR_RAISE(FieldNegFunc neg_func, CreateNegFunc(field_type)); - return std::unique_ptr(new FieldSumAgg(field_type, sum_func, neg_func)); + return std::unique_ptr(new FieldSumAgg(field_type, sum_func, neg_func, pool)); } - VariantType Agg(const VariantType& accumulator, const VariantType& input_field) override { + Result Agg(const VariantType& accumulator, + const VariantType& input_field) override { bool accumulator_null = DataDefine::IsVariantNull(accumulator); bool input_null = DataDefine::IsVariantNull(input_field); if (accumulator_null || input_null) { @@ -76,8 +78,8 @@ class FieldSumAgg : public FieldAggregator { using FieldNegFunc = std::function; FieldSumAgg(const std::shared_ptr& field_type, const FieldSumFunc& sum_func, - const FieldNegFunc& neg_func) - : FieldAggregator(std::string(NAME), field_type), + const FieldNegFunc& neg_func, const std::shared_ptr& pool) + : FieldAggregator(std::string(NAME), field_type, pool), sum_func_(sum_func), neg_func_(neg_func) {} diff --git a/src/paimon/core/mergetree/compact/aggregate/field_sum_agg_test.cpp b/src/paimon/core/mergetree/compact/aggregate/field_sum_agg_test.cpp index 994a97303..32b8dd1fb 100644 --- a/src/paimon/core/mergetree/compact/aggregate/field_sum_agg_test.cpp +++ b/src/paimon/core/mergetree/compact/aggregate/field_sum_agg_test.cpp @@ -22,14 +22,15 @@ #include "gtest/gtest.h" #include "paimon/common/utils/decimal_utils.h" #include "paimon/data/decimal.h" +#include "paimon/memory/memory_pool.h" #include "paimon/status.h" #include "paimon/testing/utils/testharness.h" namespace paimon::test { TEST(FieldSumAggTest, TestSimple) { ASSERT_OK_AND_ASSIGN(std::unique_ptr field_sum_agg, - FieldSumAgg::Create(arrow::int32())); - auto agg_ret = field_sum_agg->Agg(5, 10); + FieldSumAgg::Create(arrow::int32(), GetDefaultPool())); + auto agg_ret = field_sum_agg->Agg(5, 10).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 15); ASSERT_OK_AND_ASSIGN(auto retract_ret, field_sum_agg->Retract(5, 10)); @@ -37,17 +38,17 @@ TEST(FieldSumAggTest, TestSimple) { } TEST(FieldSumAggTest, TestNull) { ASSERT_OK_AND_ASSIGN(std::unique_ptr field_sum_agg, - FieldSumAgg::Create(arrow::int32())); + FieldSumAgg::Create(arrow::int32(), GetDefaultPool())); { - auto agg_ret = field_sum_agg->Agg(5, NullType()); + auto agg_ret = field_sum_agg->Agg(5, NullType()).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 5); } { - auto agg_ret = field_sum_agg->Agg(NullType(), 10); + auto agg_ret = field_sum_agg->Agg(NullType(), 10).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 10); } { - auto agg_ret = field_sum_agg->Agg(NullType(), NullType()); + auto agg_ret = field_sum_agg->Agg(NullType(), NullType()).value(); ASSERT_TRUE(DataDefine::IsVariantNull(agg_ret)); } @@ -67,48 +68,58 @@ TEST(FieldSumAggTest, TestNull) { TEST(FieldSumAggTest, TestVariantType) { { - ASSERT_OK_AND_ASSIGN(auto field_sum_agg, FieldSumAgg::Create(arrow::int8())); - auto agg_ret = field_sum_agg->Agg(static_cast(100), static_cast(15)); + ASSERT_OK_AND_ASSIGN(auto field_sum_agg, + FieldSumAgg::Create(arrow::int8(), GetDefaultPool())); + auto agg_ret = field_sum_agg->Agg(static_cast(100), static_cast(15)).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 115); ASSERT_OK_AND_ASSIGN(auto retract_ret, field_sum_agg->Retract(static_cast(100), static_cast(15))); ASSERT_EQ(DataDefine::GetVariantValue(retract_ret), 85); } { - ASSERT_OK_AND_ASSIGN(auto field_sum_agg, FieldSumAgg::Create(arrow::int16())); - auto agg_ret = field_sum_agg->Agg(static_cast(100), static_cast(15)); + ASSERT_OK_AND_ASSIGN(auto field_sum_agg, + FieldSumAgg::Create(arrow::int16(), GetDefaultPool())); + auto agg_ret = + field_sum_agg->Agg(static_cast(100), static_cast(15)).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 115); ASSERT_OK_AND_ASSIGN(auto retract_ret, field_sum_agg->Retract(static_cast(100), static_cast(15))); ASSERT_EQ(DataDefine::GetVariantValue(retract_ret), 85); } { - ASSERT_OK_AND_ASSIGN(auto field_sum_agg, FieldSumAgg::Create(arrow::int32())); - auto agg_ret = field_sum_agg->Agg(static_cast(100), static_cast(15)); + ASSERT_OK_AND_ASSIGN(auto field_sum_agg, + FieldSumAgg::Create(arrow::int32(), GetDefaultPool())); + auto agg_ret = + field_sum_agg->Agg(static_cast(100), static_cast(15)).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 115); ASSERT_OK_AND_ASSIGN(auto retract_ret, field_sum_agg->Retract(static_cast(100), static_cast(15))); ASSERT_EQ(DataDefine::GetVariantValue(retract_ret), 85); } { - ASSERT_OK_AND_ASSIGN(auto field_sum_agg, FieldSumAgg::Create(arrow::int64())); - auto agg_ret = field_sum_agg->Agg(static_cast(100), static_cast(15)); + ASSERT_OK_AND_ASSIGN(auto field_sum_agg, + FieldSumAgg::Create(arrow::int64(), GetDefaultPool())); + auto agg_ret = + field_sum_agg->Agg(static_cast(100), static_cast(15)).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), 115); ASSERT_OK_AND_ASSIGN(auto retract_ret, field_sum_agg->Retract(static_cast(100), static_cast(15))); ASSERT_EQ(DataDefine::GetVariantValue(retract_ret), 85); } { - ASSERT_OK_AND_ASSIGN(auto field_sum_agg, FieldSumAgg::Create(arrow::float32())); - auto agg_ret = field_sum_agg->Agg(static_cast(100.2), static_cast(15.1)); + ASSERT_OK_AND_ASSIGN(auto field_sum_agg, + FieldSumAgg::Create(arrow::float32(), GetDefaultPool())); + auto agg_ret = + field_sum_agg->Agg(static_cast(100.2), static_cast(15.1)).value(); ASSERT_NEAR(DataDefine::GetVariantValue(agg_ret), 115.3, 0.0001); ASSERT_OK_AND_ASSIGN(auto retract_ret, field_sum_agg->Retract(static_cast(100.2), static_cast(15.1))); ASSERT_NEAR(DataDefine::GetVariantValue(retract_ret), 85.1, 0.0001); } { - ASSERT_OK_AND_ASSIGN(auto field_sum_agg, FieldSumAgg::Create(arrow::float64())); - auto agg_ret = field_sum_agg->Agg(100.23, 15.11); + ASSERT_OK_AND_ASSIGN(auto field_sum_agg, + FieldSumAgg::Create(arrow::float64(), GetDefaultPool())); + auto agg_ret = field_sum_agg->Agg(100.23, 15.11).value(); ASSERT_NEAR(DataDefine::GetVariantValue(agg_ret), 115.34, 0.0001); ASSERT_OK_AND_ASSIGN(auto retract_ret, field_sum_agg->Retract(static_cast(100.23), static_cast(15.11))); @@ -119,8 +130,9 @@ TEST(FieldSumAggTest, TestVariantType) { DecimalUtils::StrToInt128("12345678998765432145678").value()); Decimal decimal2(/*precision=*/30, /*scale=*/20, DecimalUtils::StrToInt128("2345679987639475677478").value()); - ASSERT_OK_AND_ASSIGN(auto field_sum_agg, FieldSumAgg::Create(arrow::decimal128(30, 20))); - auto agg_ret = field_sum_agg->Agg(decimal1, decimal2); + ASSERT_OK_AND_ASSIGN(auto field_sum_agg, + FieldSumAgg::Create(arrow::decimal128(30, 20), GetDefaultPool())); + auto agg_ret = field_sum_agg->Agg(decimal1, decimal2).value(); ASSERT_EQ(DataDefine::GetVariantValue(agg_ret), Decimal(/*precision=*/30, /*scale=*/20, DecimalUtils::StrToInt128("14691358986404907823156").value())); @@ -132,7 +144,7 @@ TEST(FieldSumAggTest, TestVariantType) { } TEST(FieldSumAggTest, TestInvalidType) { - auto field_sum_agg = FieldSumAgg::Create(arrow::boolean()); + auto field_sum_agg = FieldSumAgg::Create(arrow::boolean(), GetDefaultPool()); ASSERT_FALSE(field_sum_agg.ok()); ASSERT_TRUE(field_sum_agg.status().ToString().find("type bool not support in FieldSumAgg") != std::string::npos) diff --git a/src/paimon/core/mergetree/compact/lookup_changelog_merge_function_wrapper_test.cpp b/src/paimon/core/mergetree/compact/lookup_changelog_merge_function_wrapper_test.cpp index cd2ebe591..0acd1b8f6 100644 --- a/src/paimon/core/mergetree/compact/lookup_changelog_merge_function_wrapper_test.cpp +++ b/src/paimon/core/mergetree/compact/lookup_changelog_merge_function_wrapper_test.cpp @@ -133,7 +133,7 @@ TEST(LookupChangelogMergeFunctionWrapperTest, TestWithLookupWithDv) { CoreOptions::FromMap({{Options::FIELDS_DEFAULT_AGG_FUNC, "sum"}})); ASSERT_OK_AND_ASSIGN(auto mfunc, AggregateMergeFunction::Create( arrow::schema({arrow::field("value", arrow::int32())}), - {"key"}, core_options)); + {"key"}, core_options, GetDefaultPool())); auto lookup_mfunc = std::make_unique(std::move(mfunc)); auto lookup = [&](const std::shared_ptr& key) -> Result> { diff --git a/src/paimon/core/mergetree/compact/lookup_merge_function_test.cpp b/src/paimon/core/mergetree/compact/lookup_merge_function_test.cpp index 54098de28..25c4c601b 100644 --- a/src/paimon/core/mergetree/compact/lookup_merge_function_test.cpp +++ b/src/paimon/core/mergetree/compact/lookup_merge_function_test.cpp @@ -40,9 +40,9 @@ TEST(LookupMergeFunctionTest, TestSimple) { auto value_schema = arrow::schema(fields); ASSERT_OK_AND_ASSIGN(CoreOptions core_options, CoreOptions::FromMap({{Options::FIELDS_DEFAULT_AGG_FUNC, "sum"}})); - ASSERT_OK_AND_ASSIGN( - std::unique_ptr agg_merge_func, - AggregateMergeFunction::Create(value_schema, /*primary_keys=*/{"k0"}, core_options)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr agg_merge_func, + AggregateMergeFunction::Create(value_schema, /*primary_keys=*/{"k0"}, + core_options, GetDefaultPool())); auto merge_func = std::make_unique(std::move(agg_merge_func)); auto pool = GetDefaultPool(); diff --git a/src/paimon/core/mergetree/compact/lookup_merge_tree_compact_rewriter_test.cpp b/src/paimon/core/mergetree/compact/lookup_merge_tree_compact_rewriter_test.cpp index 2b2678071..6459118bb 100644 --- a/src/paimon/core/mergetree/compact/lookup_merge_tree_compact_rewriter_test.cpp +++ b/src/paimon/core/mergetree/compact/lookup_merge_tree_compact_rewriter_test.cpp @@ -177,10 +177,10 @@ class LookupMergeTreeCompactRewriterTest : public ::testing::TestWithParam Result>> { - PAIMON_ASSIGN_OR_RAISE( - auto merge_func, - AggregateMergeFunction::Create( - arrow_schema_, table_schema->TrimmedPrimaryKeys().value(), options)); + PAIMON_ASSIGN_OR_RAISE(auto merge_func, + AggregateMergeFunction::Create( + arrow_schema_, table_schema->TrimmedPrimaryKeys().value(), + options, GetDefaultPool())); PAIMON_ASSIGN_OR_RAISE( std::shared_ptr> merge_function_wrapper, LookupMergeTreeCompactRewriter::CreateLookupMergeFunctionWrapper( @@ -253,10 +253,10 @@ class LookupMergeTreeCompactRewriterTest : public ::testing::TestWithParam Result>> { - PAIMON_ASSIGN_OR_RAISE( - auto merge_func, - AggregateMergeFunction::Create( - arrow_schema_, table_schema->TrimmedPrimaryKeys().value(), options)); + PAIMON_ASSIGN_OR_RAISE(auto merge_func, + AggregateMergeFunction::Create( + arrow_schema_, table_schema->TrimmedPrimaryKeys().value(), + options, GetDefaultPool())); PAIMON_ASSIGN_OR_RAISE( std::shared_ptr> merge_function_wrapper, LookupMergeTreeCompactRewriter:: diff --git a/src/paimon/core/mergetree/compact/merge_tree_compact_manager_factory.cpp b/src/paimon/core/mergetree/compact/merge_tree_compact_manager_factory.cpp index ff595fc75..eece8598c 100644 --- a/src/paimon/core/mergetree/compact/merge_tree_compact_manager_factory.cpp +++ b/src/paimon/core/mergetree/compact/merge_tree_compact_manager_factory.cpp @@ -240,12 +240,12 @@ MergeTreeCompactManagerFactory::CreateLookupRewriterWithDeletionVector( auto merge_function_wrapper_factory = [data_schema = schema_, options = options_, trimmed_primary_keys, lookup_levels_ptr = lookup_levels.get(), lookup_strategy, - dv_maintainer_ptr = dv_maintainer, + dv_maintainer_ptr = dv_maintainer, pool = pool_, user_defined_seq_comparator = user_defined_seq_comparator_]( int32_t output_level) -> Result>> { PAIMON_ASSIGN_OR_RAISE(std::unique_ptr merge_func, PrimaryKeyTableUtils::CreateMergeFunction( - data_schema, trimmed_primary_keys, options)); + data_schema, trimmed_primary_keys, options, pool)); PAIMON_ASSIGN_OR_RAISE( std::shared_ptr> merge_function_wrapper, LookupMergeTreeCompactRewriter:: @@ -274,12 +274,12 @@ MergeTreeCompactManagerFactory::CreateLookupRewriterWithDeletionVector( auto merge_function_wrapper_factory = [data_schema = schema_, options = options_, trimmed_primary_keys, lookup_levels_ptr = lookup_levels.get(), lookup_strategy, - dv_maintainer_ptr = dv_maintainer, + dv_maintainer_ptr = dv_maintainer, pool = pool_, user_defined_seq_comparator = user_defined_seq_comparator_]( int32_t output_level) -> Result>> { - PAIMON_ASSIGN_OR_RAISE( - std::unique_ptr merge_func, - PrimaryKeyTableUtils::CreateMergeFunction(data_schema, trimmed_primary_keys, options)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr merge_func, + PrimaryKeyTableUtils::CreateMergeFunction( + data_schema, trimmed_primary_keys, options, pool)); PAIMON_ASSIGN_OR_RAISE( std::shared_ptr> merge_function_wrapper, LookupMergeTreeCompactRewriter::CreateLookupMergeFunctionWrapper( @@ -317,12 +317,12 @@ MergeTreeCompactManagerFactory::CreateLookupRewriterWithoutDeletionVector( auto merge_function_wrapper_factory = [data_schema = schema_, options = options_, trimmed_primary_keys, lookup_levels_ptr = lookup_levels.get(), lookup_strategy, - dv_maintainer_ptr = dv_maintainer, + dv_maintainer_ptr = dv_maintainer, pool = pool_, user_defined_seq_comparator = user_defined_seq_comparator_]( int32_t output_level) -> Result>> { - PAIMON_ASSIGN_OR_RAISE( - std::unique_ptr merge_func, - PrimaryKeyTableUtils::CreateMergeFunction(data_schema, trimmed_primary_keys, options)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr merge_func, + PrimaryKeyTableUtils::CreateMergeFunction( + data_schema, trimmed_primary_keys, options, pool)); PAIMON_ASSIGN_OR_RAISE( std::shared_ptr> merge_function_wrapper, LookupMergeTreeCompactRewriter::CreateLookupMergeFunctionWrapper( 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..af5ab1823 100644 --- a/src/paimon/core/mergetree/compact/partial_update_merge_function.cpp +++ b/src/paimon/core/mergetree/compact/partial_update_merge_function.cpp @@ -120,13 +120,13 @@ Result> PartialUpdateMergeFunction:: const std::shared_ptr& value_schema, const std::vector& primary_keys, const CoreOptions& options, const std::map>& value_field_to_seq_group_field, - const std::set& seq_group_key_set) { + const std::set& seq_group_key_set, const std::shared_ptr& pool) { // 1. create field aggregator std::map> field_aggregators; PAIMON_ASSIGN_OR_RAISE( field_aggregators, CreateFieldAggregators(value_schema, primary_keys, options, value_field_to_seq_group_field, - seq_group_key_set)); + seq_group_key_set, pool)); // 2. create field seq comparator std::map> field_comparators; PAIMON_ASSIGN_OR_RAISE(field_comparators, @@ -223,7 +223,7 @@ PartialUpdateMergeFunction::CreateFieldAggregators( const std::shared_ptr& value_schema, const std::vector& primary_keys, const CoreOptions& options, const std::map>& value_field_to_seq_group_field, - const std::set& seq_group_key_set) { + const std::set& seq_group_key_set, const std::shared_ptr& pool) { std::map> aggregators; std::optional default_agg_func = options.GetFieldsDefaultFunc(); for (int32_t i = 0; i < value_schema->num_fields(); i++) { @@ -235,9 +235,10 @@ PartialUpdateMergeFunction::CreateFieldAggregators( } auto primary_key_iter = std::find(primary_keys.begin(), primary_keys.end(), field_name); if (primary_key_iter != primary_keys.end()) { - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr agg, - FieldAggregatorFactory::CreateFieldAggregator( - field_name, field_type, FieldPrimaryKeyAgg::NAME, options)); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr agg, + FieldAggregatorFactory::CreateFieldAggregator( + field_name, field_type, FieldPrimaryKeyAgg::NAME, options, pool)); aggregators[i] = std::move(agg); continue; } @@ -257,7 +258,7 @@ PartialUpdateMergeFunction::CreateFieldAggregators( } PAIMON_ASSIGN_OR_RAISE(std::unique_ptr agg, FieldAggregatorFactory::CreateFieldAggregator( - field_name, field_type, str_agg.value(), options)); + field_name, field_type, str_agg.value(), options, pool)); aggregators[i] = std::move(agg); } } @@ -333,16 +334,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 +351,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 +364,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->Agg(accumulator, field)); + row_->SetField(i, result); } else if (!DataDefine::IsVariantNull(field)) { row_->SetField(i, field); } @@ -383,13 +386,20 @@ void PartialUpdateMergeFunction::UpdateWithSequenceGroup(KeyValue&& kv) { } continue; } - row_->SetField(i, agg ? agg->Agg(accumulator, field) : field); + if (agg) { + PAIMON_ASSIGN_OR_RAISE(VariantType result, agg->Agg(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->AggReversed(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..8ea75755a 100644 --- a/src/paimon/core/mergetree/compact/partial_update_merge_function.h +++ b/src/paimon/core/mergetree/compact/partial_update_merge_function.h @@ -65,7 +65,7 @@ class PartialUpdateMergeFunction : public MergeFunction { const std::shared_ptr& value_schema, const std::vector& primary_keys, const CoreOptions& options, const std::map>& value_field_to_seq_group_field, - const std::set& seq_group_key_set); + const std::set& seq_group_key_set, const std::shared_ptr& pool); void Reset() override; @@ -94,7 +94,7 @@ class PartialUpdateMergeFunction : public MergeFunction { const std::shared_ptr& value_schema, const std::vector& primary_keys, const CoreOptions& options, const std::map>& value_field_to_seq_group_field, - const std::set& seq_group_key_set); + const std::set& seq_group_key_set, const std::shared_ptr& pool); bool IsEmptySequenceGroup(const KeyValue& kv, const std::shared_ptr& comparator) const; @@ -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/src/paimon/core/mergetree/compact/partial_update_merge_function_test.cpp b/src/paimon/core/mergetree/compact/partial_update_merge_function_test.cpp index a403af59b..838542c55 100644 --- a/src/paimon/core/mergetree/compact/partial_update_merge_function_test.cpp +++ b/src/paimon/core/mergetree/compact/partial_update_merge_function_test.cpp @@ -17,9 +17,12 @@ #include "paimon/core/mergetree/compact/partial_update_merge_function.h" #include +#include #include #include +#include "DataSketches/hll.hpp" +#include "arrow/api.h" #include "arrow/type.h" #include "gtest/gtest.h" #include "paimon/common/data/data_define.h" @@ -31,6 +34,7 @@ #include "paimon/core/mergetree/compact/aggregate/field_sum_agg.h" #include "paimon/core/schema/table_schema.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" @@ -83,8 +87,8 @@ class PartialUpdateMergeFunctionTest : public testing::Test { PAIMON_RETURN_NOT_OK(PartialUpdateMergeFunction::ParseSequenceGroupFields( options, &value_field_to_seq_group_field, &seq_group_key_set)); return PartialUpdateMergeFunction::Create(value_schema, /*primary_keys=*/{"f0"}, options, - value_field_to_seq_group_field, - seq_group_key_set); + value_field_to_seq_group_field, seq_group_key_set, + GetDefaultPool()); } std::string CreateMergeFunctionWithInvalidOptions( @@ -118,10 +122,10 @@ class PartialUpdateMergeFunctionTest : public testing::Test { table_schema, value_field_to_seq_group_field, ©_value_fields)); EXPECT_EQ(copy_value_fields, expected_completed_value_fields); auto value_schema = DataField::ConvertDataFieldsToArrowSchema(copy_value_fields); - EXPECT_OK_AND_ASSIGN( - std::unique_ptr mfunc, - PartialUpdateMergeFunction::Create(value_schema, {"f0"}, options, - value_field_to_seq_group_field, seq_group_key_set)); + EXPECT_OK_AND_ASSIGN(std::unique_ptr mfunc, + PartialUpdateMergeFunction::Create( + value_schema, {"f0"}, options, value_field_to_seq_group_field, + seq_group_key_set, GetDefaultPool())); return mfunc; } @@ -436,10 +440,10 @@ TEST_F(PartialUpdateMergeFunctionTest, TestAdjustProjectionCreateDirectly) { ASSERT_OK(PartialUpdateMergeFunction::ParseSequenceGroupFields( options, &value_field_to_seq_group_field, &seq_group_key_set)); - ASSERT_NOK_WITH_MSG( - PartialUpdateMergeFunction::Create(value_schema, {"f0"}, options, - value_field_to_seq_group_field, seq_group_key_set), - "cannot find sequence group field f4 in value schema, unexpected."); + ASSERT_NOK_WITH_MSG(PartialUpdateMergeFunction::Create(value_schema, {"f0"}, options, + value_field_to_seq_group_field, + seq_group_key_set, GetDefaultPool()), + "cannot find sequence group field f4 in value schema, unexpected."); } TEST_F(PartialUpdateMergeFunctionTest, TestFirstValue) { @@ -750,7 +754,7 @@ TEST_F(PartialUpdateMergeFunctionTest, TestCreateFieldAggregatorsWithDefaultAgg) ASSERT_OK_AND_ASSIGN(aggs, PartialUpdateMergeFunction::CreateFieldAggregators( value_schema, /*primary_keys=*/{"p0"}, options, value_field_to_seq_group_field, - seq_group_key_set)); + seq_group_key_set, GetDefaultPool())); ASSERT_EQ(4, aggs.size()); // test primary key: p0 ASSERT_TRUE(dynamic_cast(aggs[0].get())); @@ -781,7 +785,7 @@ TEST_F(PartialUpdateMergeFunctionTest, TestCreateFieldAggregatorsWithoutDefaultA ASSERT_OK_AND_ASSIGN(aggs, PartialUpdateMergeFunction::CreateFieldAggregators( value_schema, /*primary_keys=*/{"p0"}, options, value_field_to_seq_group_field, - seq_group_key_set)); + seq_group_key_set, GetDefaultPool())); ASSERT_EQ(2, aggs.size()); // test primary key: p0 ASSERT_TRUE(dynamic_cast(aggs[0].get())); @@ -860,9 +864,10 @@ TEST_F(PartialUpdateMergeFunctionTest, TestInitRowWithNullableFieldOnDelete) { std::set seq_group_key_set; ASSERT_OK(PartialUpdateMergeFunction::ParseSequenceGroupFields( options, &value_field_to_seq_group_field, &seq_group_key_set)); - ASSERT_OK_AND_ASSIGN(auto mfunc, PartialUpdateMergeFunction::Create( - value_schema, /*primary_keys=*/{"f0"}, options, - value_field_to_seq_group_field, seq_group_key_set)); + ASSERT_OK_AND_ASSIGN(auto mfunc, + PartialUpdateMergeFunction::Create(value_schema, /*primary_keys=*/{"f0"}, + options, value_field_to_seq_group_field, + seq_group_key_set, GetDefaultPool())); mfunc->Reset(); // insert some data first @@ -872,4 +877,60 @@ TEST_F(PartialUpdateMergeFunctionTest, TestInitRowWithNullableFieldOnDelete) { // after delete with removeRecordOnDelete, row is re-initialized via initRow CheckResult(mfunc, {1, 2, 2, NullType()}); } + +namespace { + +std::shared_ptr HllBytes(std::initializer_list values) { + datasketches::hll_sketch sketch(12, datasketches::HLL_4); + for (int32_t value : values) { + sketch.update(value); + } + std::vector serialized = sketch.serialize_compact(); + std::shared_ptr bytes = Bytes::AllocateBytes(serialized.size(), GetDefaultPool().get()); + std::memcpy(bytes->data(), serialized.data(), bytes->size()); + return bytes; +} + +// 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 + +// An out-of-order record routes through AggReversed, which swaps the arguments +TEST_F(PartialUpdateMergeFunctionTest, SequenceGroupKeepsBinaryAccumulatorOnOlderNullField) { + std::vector data_fields = {DataField(0, arrow::field("f0", arrow::int32())), + DataField(1, arrow::field("f1", arrow::binary())), + DataField(2, arrow::field("f2", arrow::int32()))}; + auto value_schema = DataField::ConvertDataFieldsToArrowSchema(data_fields); + ASSERT_OK_AND_ASSIGN(CoreOptions options, + CoreOptions::FromMap({{"fields.f2.sequence-group", "f1"}, + {"fields.f1.aggregate-function", "hll_sketch"}})); + std::map> value_field_to_seq_group_field; + std::set seq_group_key_set; + ASSERT_OK(PartialUpdateMergeFunction::ParseSequenceGroupFields( + options, &value_field_to_seq_group_field, &seq_group_key_set)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr mfunc, + PartialUpdateMergeFunction::Create(value_schema, /*primary_keys=*/{"f0"}, + options, value_field_to_seq_group_field, + seq_group_key_set, GetDefaultPool())); + mfunc->Reset(); + Add(mfunc, {1, HllBytes({1, 2, 3}), 10}); + // newer record, so the forward path unions and stores an owned buffer in the row + Add(mfunc, {1, HllBytes({3, 4, 5}), 20}); + // older record with a null sketch, which takes the reversed null path + Add(mfunc, {1, NullType(), 15}); + + std::vector> scribbled = ScribbleFreedMemory(); + std::string_view bytes = mfunc->row_->GetStringView(1); + ASSERT_NEAR( + 5.0, datasketches::hll_sketch::deserialize(bytes.data(), bytes.size()).get_estimate(), 0.1); +} } // namespace paimon::test diff --git a/src/paimon/core/mergetree/compact/sort_merge_reader_test.cpp b/src/paimon/core/mergetree/compact/sort_merge_reader_test.cpp index 31d991449..6a98ab23d 100644 --- a/src/paimon/core/mergetree/compact/sort_merge_reader_test.cpp +++ b/src/paimon/core/mergetree/compact/sort_merge_reader_test.cpp @@ -150,9 +150,9 @@ class SortMergeReaderTest : public testing::Test { const std::vector& primary_keys, const CoreOptions& core_options, const std::vector& expected) const { for (auto batch_size : {1, 2, 3, 4, 100}) { - ASSERT_OK_AND_ASSIGN( - std::unique_ptr mfunc, - AggregateMergeFunction::Create(value_schema, primary_keys, core_options)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr mfunc, + AggregateMergeFunction::Create(value_schema, primary_keys, + core_options, GetDefaultPool())); auto merge_function_wrapper = std::make_shared(std::move(mfunc)); std::vector> merged_readers; diff --git a/src/paimon/core/operation/file_store_write.cpp b/src/paimon/core/operation/file_store_write.cpp index ceb27b073..56fcbe63f 100644 --- a/src/paimon/core/operation/file_store_write.cpp +++ b/src/paimon/core/operation/file_store_write.cpp @@ -180,9 +180,9 @@ Result> FileStoreWrite::Create(std::unique_ptrPrimaryKeys(); - PAIMON_ASSIGN_OR_RAISE( - std::unique_ptr merge_function, - PrimaryKeyTableUtils::CreateMergeFunction(arrow_schema, primary_keys, options)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr merge_function, + PrimaryKeyTableUtils::CreateMergeFunction( + arrow_schema, primary_keys, options, ctx->GetMemoryPool())); if (options.NeedLookup() && options.GetMergeEngine() != MergeEngine::FIRST_ROW) { // don't wrap first row, it is already OK merge_function = std::make_unique(std::move(merge_function)); diff --git a/src/paimon/core/operation/merge_file_split_read.cpp b/src/paimon/core/operation/merge_file_split_read.cpp index e1cee5a08..383e776a5 100644 --- a/src/paimon/core/operation/merge_file_split_read.cpp +++ b/src/paimon/core/operation/merge_file_split_read.cpp @@ -170,7 +170,7 @@ MergeFileSplitRead::GetMergeFunctionWrapper() { // raise errors when creating MergeFileSplitRead at the beginning. PAIMON_ASSIGN_OR_RAISE( merge_function_wrapper_, - CreateMergeFunctionWrapper(options_, context_->GetTableSchema(), value_schema_)); + CreateMergeFunctionWrapper(options_, context_->GetTableSchema(), value_schema_, pool_)); } return merge_function_wrapper_; } @@ -178,10 +178,11 @@ MergeFileSplitRead::GetMergeFunctionWrapper() { Result>> MergeFileSplitRead::CreateMergeFunctionWrapper(const CoreOptions& core_options, const std::shared_ptr& table_schema, - const std::shared_ptr& value_schema) { + const std::shared_ptr& value_schema, + const std::shared_ptr& pool) { PAIMON_ASSIGN_OR_RAISE(std::unique_ptr merge_function, PrimaryKeyTableUtils::CreateMergeFunction( - value_schema, table_schema->PrimaryKeys(), core_options)); + value_schema, table_schema->PrimaryKeys(), core_options, pool)); if (core_options.NeedLookup() && core_options.GetMergeEngine() != MergeEngine::FIRST_ROW) { // don't wrap first row, it is already OK merge_function = std::make_unique(std::move(merge_function)); diff --git a/src/paimon/core/operation/merge_file_split_read.h b/src/paimon/core/operation/merge_file_split_read.h index 170b2f50a..1f94878ac 100644 --- a/src/paimon/core/operation/merge_file_split_read.h +++ b/src/paimon/core/operation/merge_file_split_read.h @@ -157,7 +157,8 @@ class MergeFileSplitRead : public AbstractSplitRead { static Result>> CreateMergeFunctionWrapper( const CoreOptions& core_options, const std::shared_ptr& table_schema, - const std::shared_ptr& value_schema); + const std::shared_ptr& value_schema, + const std::shared_ptr& pool); static Status GenerateKeyValueReadSchema( const TableSchema& table_schema, const CoreOptions& options, diff --git a/src/paimon/core/utils/primary_key_table_utils.cpp b/src/paimon/core/utils/primary_key_table_utils.cpp index 168f4c685..a9f18b7e0 100644 --- a/src/paimon/core/utils/primary_key_table_utils.cpp +++ b/src/paimon/core/utils/primary_key_table_utils.cpp @@ -43,7 +43,8 @@ namespace paimon { Result> PrimaryKeyTableUtils::CreateMergeFunction( const std::shared_ptr& value_schema, - const std::vector& primary_keys, const CoreOptions& options) { + const std::vector& primary_keys, const CoreOptions& options, + const std::shared_ptr& pool) { auto merge_engine = options.GetMergeEngine(); std::unique_ptr merge_function; if (merge_engine == MergeEngine::DEDUPLICATE) { @@ -51,17 +52,17 @@ Result> PrimaryKeyTableUtils::CreateMergeFunction } else if (merge_engine == MergeEngine::FIRST_ROW) { merge_function = std::make_unique(options.IgnoreDelete()); } else if (merge_engine == MergeEngine::AGGREGATE) { - PAIMON_ASSIGN_OR_RAISE(merge_function, - AggregateMergeFunction::Create(value_schema, primary_keys, options)); + PAIMON_ASSIGN_OR_RAISE(merge_function, AggregateMergeFunction::Create( + value_schema, primary_keys, options, pool)); } else if (merge_engine == MergeEngine::PARTIAL_UPDATE) { std::map> value_field_to_seq_group_field; std::set seq_group_key_set; PAIMON_RETURN_NOT_OK(PartialUpdateMergeFunction::ParseSequenceGroupFields( options, &value_field_to_seq_group_field, &seq_group_key_set)); PAIMON_ASSIGN_OR_RAISE( - merge_function, - PartialUpdateMergeFunction::Create(value_schema, primary_keys, options, - value_field_to_seq_group_field, seq_group_key_set)); + merge_function, PartialUpdateMergeFunction::Create(value_schema, primary_keys, options, + value_field_to_seq_group_field, + seq_group_key_set, pool)); } else { return Status::Invalid( "only support deduplicate/partial-update/aggregation/first-row merge engine."); diff --git a/src/paimon/core/utils/primary_key_table_utils.h b/src/paimon/core/utils/primary_key_table_utils.h index d8413808f..523c18e0e 100644 --- a/src/paimon/core/utils/primary_key_table_utils.h +++ b/src/paimon/core/utils/primary_key_table_utils.h @@ -31,6 +31,7 @@ namespace paimon { class MergeFunction; class CoreOptions; +class MemoryPool; class FieldsComparator; class DataField; @@ -39,9 +40,18 @@ class PrimaryKeyTableUtils { PrimaryKeyTableUtils() = delete; ~PrimaryKeyTableUtils() = delete; + /// Create the merge engine configured for a primary key table. + /// + /// @param value_schema Schema of the value part of a KeyValue. + /// @param primary_keys Primary key field names. + /// @param options Table options selecting the merge engine. + /// @param pool Pool the merge engine charges its allocations to. Aggregating merge engines + /// allocate per merged value, so the caller's pool must be threaded through. + /// @return The merge function, or an error Status for an unsupported merge engine. static Result> CreateMergeFunction( const std::shared_ptr& value_schema, - const std::vector& primary_keys, const CoreOptions& options); + const std::vector& primary_keys, const CoreOptions& options, + const std::shared_ptr& pool); static Result> CreateSequenceFieldsComparator( const std::vector& value_fields, const CoreOptions& options); diff --git a/src/paimon/core/utils/primary_key_table_utils_test.cpp b/src/paimon/core/utils/primary_key_table_utils_test.cpp index 8b35dd485..6027e17b0 100644 --- a/src/paimon/core/utils/primary_key_table_utils_test.cpp +++ b/src/paimon/core/utils/primary_key_table_utils_test.cpp @@ -73,9 +73,9 @@ TEST(PrimaryKeyTableUtilsTest, TestCreateFirstRowMergeFunctionWithIgnoreDelete) ASSERT_OK_AND_ASSIGN(CoreOptions core_options, CoreOptions::FromMap({{Options::MERGE_ENGINE, "first-row"}, {ignore_delete_key, "true"}})); - ASSERT_OK_AND_ASSIGN( - std::unique_ptr merge_function, - PrimaryKeyTableUtils::CreateMergeFunction(value_schema, {"k0"}, core_options)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr merge_function, + PrimaryKeyTableUtils::CreateMergeFunction( + value_schema, {"k0"}, core_options, GetDefaultPool())); merge_function->Reset(); KeyValue insert_kv(RowKind::Insert(), /*sequence_number=*/0, /*level=*/0, /*key=*/ @@ -99,9 +99,9 @@ TEST(PrimaryKeyTableUtilsTest, TestCreateFirstRowMergeFunctionWithIgnoreDelete) // Without the option, the first-row merge engine still rejects retract records. ASSERT_OK_AND_ASSIGN(CoreOptions core_options, CoreOptions::FromMap({{Options::MERGE_ENGINE, "first-row"}})); - ASSERT_OK_AND_ASSIGN( - std::unique_ptr merge_function, - PrimaryKeyTableUtils::CreateMergeFunction(value_schema, {"k0"}, core_options)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr merge_function, + PrimaryKeyTableUtils::CreateMergeFunction(value_schema, {"k0"}, + core_options, GetDefaultPool())); merge_function->Reset(); KeyValue delete_kv(RowKind::Delete(), /*sequence_number=*/0, /*level=*/0, /*key=*/ BinaryRowGenerator::GenerateRowPtr({10}, pool.get()), diff --git a/test/inte/CMakeLists.txt b/test/inte/CMakeLists.txt index 58b002764..c9d08195d 100644 --- a/test/inte/CMakeLists.txt +++ b/test/inte/CMakeLists.txt @@ -102,6 +102,7 @@ if(PAIMON_BUILD_TESTS) paimon_shared ${TEST_STATIC_LINK_LIBS} test_utils_static + DataSketches ${GTEST_LINK_TOOLCHAIN}) add_paimon_test(nested_column_pruning_inte_test diff --git a/test/inte/pk_compaction_inte_test.cpp b/test/inte/pk_compaction_inte_test.cpp index 47a915021..4c97c5fba 100644 --- a/test/inte/pk_compaction_inte_test.cpp +++ b/test/inte/pk_compaction_inte_test.cpp @@ -16,10 +16,13 @@ #include #include +#include #include #include #include +#include "DataSketches/hll.hpp" +#include "DataSketches/theta_sketch.hpp" #include "arrow/c/bridge.h" #include "arrow/ipc/json_simple.h" #include "arrow/type.h" @@ -293,6 +296,33 @@ class PkCompactionInteTest : public ::testing::Test, } } + // Read every row of the table, for fields whose expected value cannot be spelled out as JSON. + // `consume` runs while the reader is still alive, because the arrow arrays are allocated from + // a pool the reader owns and must not outlive it. + template + void ScanAllRows(const std::string& table_path, Fn consume) { + std::map options = {{Options::FILE_SYSTEM, "local"}}; + ScanContextBuilder scan_context_builder(table_path); + scan_context_builder.WithStreamingMode(false).SetOptions(options).AddOption( + Options::SCAN_MODE, StartupMode::LatestFull().ToString()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr scan_context, + scan_context_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr table_scan, + TableScan::Create(std::move(scan_context))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result_plan, table_scan->CreatePlan()); + ReadContextBuilder read_context_builder(table_path); + read_context_builder.SetOptions(options); + ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, + read_context_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr table_read, + TableRead::Create(std::move(read_context))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch_reader, + table_read->CreateReader(result_plan->Splits())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, + ReadResultCollector::CollectResult(batch_reader.get())); + consume(result); + } + // Helper: check whether compact commit messages contain new DV index files. bool HasDeletionVectorIndexFiles( const std::vector>& commit_messages) { @@ -3316,6 +3346,194 @@ TEST_F(PkCompactionInteTest, PkDvAndAggWithIOException) { ASSERT_TRUE(run_complete); } +// End-to-end coverage for the collect / merge_map / nested_update aggregators: values are merged +// both on read (across level-0 files) and during compaction, then verified from the written files. +TEST_F(PkCompactionInteTest, AggCollectMergeMapAndNestedUpdate) { + auto map_type = + std::make_shared(arrow::field("key", arrow::int32(), /*nullable=*/false), + arrow::field("value", arrow::int32())); + arrow::FieldVector fields = { + arrow::field("f0", arrow::utf8()), // PK + arrow::field("f1", arrow::list(arrow::int32())), arrow::field("f2", map_type), + arrow::field("f3", arrow::list(arrow::struct_({arrow::field("id", arrow::int32()), + arrow::field("name", arrow::utf8())})))}; + std::map options = {{Options::FILE_FORMAT, "parquet"}, + {Options::BUCKET, "1"}, + {Options::FILE_SYSTEM, "local"}, + {Options::MERGE_ENGINE, "aggregation"}, + {Options::DELETION_VECTORS_ENABLED, "true"}, + {"fields.f1.aggregate-function", "collect"}, + {"fields.f1.distinct", "true"}, + {"fields.f2.aggregate-function", "merge_map"}, + {"fields.f3.aggregate-function", "nested_update"}, + {"fields.f3.nested-key", "id"}}; + CreateTable(fields, /*partition_keys=*/{}, /*primary_keys=*/{"f0"}, options); + std::string table_path = TablePath(); + auto data_type = arrow::struct_(fields); + int64_t commit_id = 0; + + // Step 1: initial batch, then full compact so it lands in the max level. + { + auto array = arrow::ipc::internal::json::ArrayFromJSON(data_type, R"([ + ["Alice", [1,2], [[1,10]], [[1,"a"]]], + ["Bob", [5], [[5,50]], [[5,"e"]]] + ])") + .ValueOrDie(); + ASSERT_OK(WriteAndCommit(table_path, {}, 0, array, commit_id++)); + } + ASSERT_OK_AND_ASSIGN( + [[maybe_unused]] auto upgrade_msgs, + CompactAndCommit(table_path, {}, 0, /*full_compaction=*/true, commit_id++)); + + // Step 2: two more level-0 files overlapping the max-level rows. + { + auto array = arrow::ipc::internal::json::ArrayFromJSON(data_type, R"([ + ["Alice", [2,3], [[1,11],[2,20]], [[1,"A"],[2,"b"]]] + ])") + .ValueOrDie(); + ASSERT_OK(WriteAndCommit(table_path, {}, 0, array, commit_id++)); + } + { + auto array = arrow::ipc::internal::json::ArrayFromJSON(data_type, R"([ + ["Alice", [3,4], [[3,30]], [[2,"B"]]], + ["Bob", [6], [[6,60]], [[6,"f"]]] + ])") + .ValueOrDie(); + ASSERT_OK(WriteAndCommit(table_path, {}, 0, array, commit_id++)); + } + + // Step 3: non-full compact merges the level-0 files against the max-level file. + ASSERT_OK_AND_ASSIGN( + [[maybe_unused]] auto dv_compact_msgs, + CompactAndCommit(table_path, {}, 0, /*full_compaction=*/false, commit_id++)); + + // collect keeps insertion order and drops the duplicated 2 and 3, + // merge_map overwrites key 1 in place, nested_update replaces key 1 and then key 2. + const std::string expected = R"([ + [0, "Alice", [1,2,3,4], [[1,11],[2,20],[3,30]], [[1,"A"],[2,"B"]]], + [0, "Bob", [5,6], [[5,50],[6,60]], [[5,"e"],[6,"f"]]] + ])"; + { + std::map, std::string> expected_data; + expected_data[std::make_pair("", 0)] = expected; + ScanAndVerify(table_path, fields, expected_data); + } + + // Step 4: the same values must survive a full compaction and come back from the files. + ASSERT_OK_AND_ASSIGN( + [[maybe_unused]] auto full_compact_msgs, + CompactAndCommit(table_path, {}, 0, /*full_compaction=*/true, commit_id++)); + { + std::map, std::string> expected_data; + expected_data[std::make_pair("", 0)] = expected; + ScanAndVerify(table_path, fields, expected_data); + } +} + +// End-to-end coverage for the sketch aggregators. The last batch writes null sketches, which is the +// path where the aggregator must hand back an owning copy of the accumulator rather than a view +// into the row buffer that is about to be overwritten. +TEST_F(PkCompactionInteTest, AggHllAndThetaSketches) { + arrow::FieldVector fields = {arrow::field("f0", arrow::utf8()), // PK + arrow::field("f1", arrow::binary()), + arrow::field("f2", arrow::binary())}; + std::map options = {{Options::FILE_FORMAT, "parquet"}, + {Options::BUCKET, "1"}, + {Options::FILE_SYSTEM, "local"}, + {Options::MERGE_ENGINE, "aggregation"}, + {Options::DELETION_VECTORS_ENABLED, "true"}, + {"fields.f1.aggregate-function", "hll_sketch"}, + {"fields.f2.aggregate-function", "theta_sketch"}}; + CreateTable(fields, /*partition_keys=*/{}, /*primary_keys=*/{"f0"}, options); + std::string table_path = TablePath(); + int64_t commit_id = 0; + + auto hll = [](std::initializer_list values) { + datasketches::hll_sketch sketch(12, datasketches::HLL_4); + for (int32_t value : values) { + sketch.update(value); + } + return sketch.serialize_compact(); + }; + auto 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 sketch.compact(/*ordered=*/true).serialize(); + }; + // Sketches are opaque binaries, so build the arrays directly instead of via JSON. + auto make_batch = [&fields](const std::vector& keys, + const std::vector>>& hlls, + const std::vector>>& thetas) { + arrow::StringBuilder key_builder; + EXPECT_TRUE(key_builder.AppendValues(keys).ok()); + arrow::ArrayVector children(3); + children[0] = key_builder.Finish().ValueOrDie(); + const std::vector>>* columns[] = {&hlls, &thetas}; + for (int32_t column = 0; column < 2; ++column) { + arrow::BinaryBuilder builder; + for (const std::optional>& value : *columns[column]) { + if (value.has_value()) { + EXPECT_TRUE(builder.Append(value->data(), value->size()).ok()); + } else { + EXPECT_TRUE(builder.AppendNull().ok()); + } + } + children[column + 1] = builder.Finish().ValueOrDie(); + } + return std::static_pointer_cast( + arrow::StructArray::Make(children, fields).ValueOrDie()); + }; + + ASSERT_OK(WriteAndCommit( + table_path, {}, 0, + make_batch({"Alice", "Bob"}, {hll({1, 2, 3}), hll({7})}, {theta({1, 2, 3}), theta({7})}), + commit_id++)); + ASSERT_OK_AND_ASSIGN( + [[maybe_unused]] auto upgrade_msgs, + CompactAndCommit(table_path, {}, 0, /*full_compaction=*/true, commit_id++)); + + ASSERT_OK(WriteAndCommit(table_path, {}, 0, + make_batch({"Alice"}, {hll({3, 4, 5})}, {theta({3, 4, 5})}), + commit_id++)); + // null on both sides of the union, so the accumulator has to be copied out of the row + ASSERT_OK(WriteAndCommit(table_path, {}, 0, + make_batch({"Alice"}, {std::nullopt}, {std::nullopt}), commit_id++)); + ASSERT_OK_AND_ASSIGN( + [[maybe_unused]] auto full_compact_msgs, + CompactAndCommit(table_path, {}, 0, /*full_compaction=*/true, commit_id++)); + + std::map> estimates; + ScanAllRows(table_path, [&estimates](const std::shared_ptr& result) { + for (const std::shared_ptr& chunk : result->chunks()) { + auto rows = std::static_pointer_cast(chunk); + auto keys = std::static_pointer_cast(rows->field(1)); + auto hll_column = std::static_pointer_cast(rows->field(2)); + auto theta_column = std::static_pointer_cast(rows->field(3)); + for (int64_t i = 0; i < rows->length(); ++i) { + ASSERT_FALSE(hll_column->IsNull(i)); + ASSERT_FALSE(theta_column->IsNull(i)); + std::string_view hll_bytes = hll_column->GetView(i); + std::string_view theta_bytes = theta_column->GetView(i); + estimates[keys->GetString(i)] = { + datasketches::hll_sketch::deserialize(hll_bytes.data(), hll_bytes.size()) + .get_estimate(), + datasketches::compact_theta_sketch::deserialize(theta_bytes.data(), + theta_bytes.size()) + .get_estimate()}; + } + } + }); + + ASSERT_EQ(2, estimates.size()); + ASSERT_NEAR(5.0, estimates["Alice"].first, 0.1); + ASSERT_DOUBLE_EQ(5.0, estimates["Alice"].second); + ASSERT_NEAR(1.0, estimates["Bob"].first, 0.1); + ASSERT_DOUBLE_EQ(1.0, estimates["Bob"].second); +} + std::vector GetTestValuesForCompactionInteTest() { std::vector values; values.emplace_back("parquet"); 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"