feat(core): Add bucket function implementation (alibaba#218)

Author: ChaomingZhangCN

include/paimon/defs.h

@@ -414,6 +414,9 @@ struct PAIMON_EXPORT Options {
     /// "lookup.cache.high-priority-pool-ratio" - The fraction of cache memory that is reserved for
     /// high-priority data like index, filter. Default value is 0.25.
     static const char LOOKUP_CACHE_HIGH_PRIO_POOL_RATIO[];
+    /// "bucket-function.type" - The bucket function type for paimon bucket.
+    /// Values can be: "default", "mod", "hive". Default value is "default".
+    static const char BUCKET_FUNCTION_TYPE[];
     /// "lookup.cache-file-retention" - The cached files retention time for lookup.
     /// After the file expires, if there is a need for access, it will be re-read from the DFS
     /// to build an index on the local disk. Default value is 1 hour.

include/paimon/utils/bucket_function_type.h

@@ -0,0 +1,50 @@
+/*
+ * 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 <cstdint>
+
+#include "paimon/defs.h"
+#include "paimon/visibility.h"
+
+namespace paimon {
+
+/// Specifies the bucket function type for paimon bucket.
+/// This determines how rows are assigned to buckets during data writing.
+enum class BucketFunctionType {
+    /// The default bucket function which will use arithmetic:
+    /// bucket_id = abs(hash_bucket_binary_row % numBuckets) to get bucket.
+    DEFAULT = 1,
+    /// The modulus bucket function which will use modulus arithmetic:
+    /// bucket_id = floorMod(bucket_key_value, numBuckets) to get bucket.
+    /// Note: the bucket key must be a single field of INT or BIGINT datatype.
+    MOD = 2,
+    /// The hive bucket function which will use hive-compatible hash arithmetic to get bucket.
+    HIVE = 3
+};
+
+/// Describes a field's type information needed for Hive hashing.
+struct PAIMON_EXPORT HiveFieldInfo {
+    FieldType type;
+    int32_t precision = 0;  // Used for DECIMAL type
+    int32_t scale = 0;      // Used for DECIMAL type
+
+    explicit HiveFieldInfo(FieldType t) : type(t) {}
+    HiveFieldInfo(FieldType t, int32_t p, int32_t s) : type(t), precision(p), scale(s) {}
+};
+
+}  // namespace paimon

include/paimon/utils/bucket_id_calculator.h

@@ -17,16 +17,19 @@
 #pragma once
 #include <cstdint>
 #include <memory>
+#include <vector>
 
 #include "paimon/memory/memory_pool.h"
 #include "paimon/result.h"
 #include "paimon/status.h"
+#include "paimon/utils/bucket_function_type.h"
 #include "paimon/visibility.h"
 
 struct ArrowSchema;
 struct ArrowArray;
 
 namespace paimon {
+class BucketFunction;
 class MemoryPool;
 
 /// Calculator for determining bucket ids based on the given bucket keys.
@@ -35,18 +38,40 @@ class MemoryPool;
 /// hash-based distribution to ensure even data distribution across buckets.
 class PAIMON_EXPORT BucketIdCalculator {
  public:
-    /// Create `BucketIdCalculator` with custom memory pool.
+    /// Create `BucketIdCalculator` with default bucket function.
     /// @param is_pk_table Whether this is for a primary key table.
     /// @param num_buckets Number of buckets.
     /// @param pool Memory pool for memory allocation.
     static Result<std::unique_ptr<BucketIdCalculator>> Create(
         bool is_pk_table, int32_t num_buckets, const std::shared_ptr<MemoryPool>& pool);
 
-    /// Create `BucketIdCalculator` with default memory pool.
+    /// Create `BucketIdCalculator` with a custom bucket function.
     /// @param is_pk_table Whether this is for a primary key table.
     /// @param num_buckets Number of buckets.
-    static Result<std::unique_ptr<BucketIdCalculator>> Create(bool is_pk_table,
-                                                              int32_t num_buckets);
+    /// @param bucket_function The bucket function to use for bucket assignment.
+    /// @param pool Memory pool for memory allocation.
+    static Result<std::unique_ptr<BucketIdCalculator>> Create(
+        bool is_pk_table, int32_t num_buckets, std::unique_ptr<BucketFunction> bucket_function,
+        const std::shared_ptr<MemoryPool>& pool);
+
+    /// Create `BucketIdCalculator` with MOD bucket function.
+    /// @param is_pk_table Whether this is for a primary key table.
+    /// @param num_buckets Number of buckets.
+    /// @param bucket_key_type The type of the single bucket key field. Must be INT or BIGINT.
+    /// @param pool Memory pool for memory allocation.
+    static Result<std::unique_ptr<BucketIdCalculator>> CreateMod(
+        bool is_pk_table, int32_t num_buckets, FieldType bucket_key_type,
+        const std::shared_ptr<MemoryPool>& pool);
+
+    /// Create `BucketIdCalculator` with HIVE bucket function.
+    /// @param is_pk_table Whether this is for a primary key table.
+    /// @param num_buckets Number of buckets.
+    /// @param field_infos The detailed type info of all fields in the bucket key row.
+    /// @param pool Memory pool for memory allocation.
+    static Result<std::unique_ptr<BucketIdCalculator>> CreateHive(
+        bool is_pk_table, int32_t num_buckets, const std::vector<HiveFieldInfo>& field_infos,
+        const std::shared_ptr<MemoryPool>& pool);
+
     /// Calculate bucket ids for the given bucket keys.
     /// @param bucket_keys Arrow struct array containing the bucket key values.
     /// @param bucket_schema Arrow schema describing the structure of bucket_keys.
@@ -57,12 +82,16 @@ class PAIMON_EXPORT BucketIdCalculator {
     Status CalculateBucketIds(ArrowArray* bucket_keys, ArrowSchema* bucket_schema,
                               int32_t* bucket_ids) const;
 
+    /// Destructor
+    ~BucketIdCalculator();
+
  private:
-    BucketIdCalculator(int32_t num_buckets, const std::shared_ptr<MemoryPool>& pool)
-        : num_buckets_(num_buckets), pool_(pool) {}
+    BucketIdCalculator(int32_t num_buckets, std::unique_ptr<BucketFunction> bucket_function,
+                       const std::shared_ptr<MemoryPool>& pool);
 
  private:
     int32_t num_buckets_;
+    std::unique_ptr<BucketFunction> bucket_function_;
     std::shared_ptr<MemoryPool> pool_;
 };
 }  // namespace paimon

src/paimon/CMakeLists.txt

@@ -149,6 +149,8 @@ set(PAIMON_CORE_SRCS
     core/disk/io_manager.cpp
     core/append/append_only_writer.cpp
     core/append/bucketed_append_compact_manager.cpp
+    core/bucket/hive_bucket_function.cpp
+    core/bucket/mod_bucket_function.cpp
     core/casting/binary_to_string_cast_executor.cpp
     core/casting/boolean_to_decimal_cast_executor.cpp
     core/casting/boolean_to_numeric_cast_executor.cpp
@@ -514,6 +516,9 @@ if(PAIMON_BUILD_TESTS)
                     SOURCES
                     core/append/append_only_writer_test.cpp
                     core/append/bucketed_append_compact_manager_test.cpp
+                    core/bucket/default_bucket_function_test.cpp
+                    core/bucket/hive_bucket_function_test.cpp
+                    core/bucket/mod_bucket_function_test.cpp
                     core/casting/cast_executor_factory_test.cpp
                     core/casting/cast_executor_test.cpp
                     core/casting/casted_row_test.cpp

src/paimon/common/defs.cpp

@@ -117,6 +117,7 @@ const char Options::LOOKUP_COMPACT[] = "lookup-compact";
 const char Options::LOOKUP_COMPACT_MAX_INTERVAL[] = "lookup-compact.max-interval";
 const char Options::LOOKUP_CACHE_MAX_MEMORY_SIZE[] = "lookup.cache-max-memory-size";
 const char Options::LOOKUP_CACHE_HIGH_PRIO_POOL_RATIO[] = "lookup.cache.high-priority-pool-ratio";
+const char Options::BUCKET_FUNCTION_TYPE[] = "bucket-function.type";
 const char Options::LOOKUP_CACHE_FILE_RETENTION[] = "lookup.cache-file-retention";
 const char Options::LOOKUP_CACHE_MAX_DISK_SIZE[] = "lookup.cache-max-disk-size";
 

src/paimon/common/utils/bucket_id_calculator.cpp

@@ -44,6 +44,10 @@
 #include "paimon/common/utils/arrow/status_utils.h"
 #include "paimon/common/utils/date_time_utils.h"
 #include "paimon/common/utils/scope_guard.h"
+#include "paimon/core/bucket/bucket_function.h"
+#include "paimon/core/bucket/default_bucket_function.h"
+#include "paimon/core/bucket/hive_bucket_function.h"
+#include "paimon/core/bucket/mod_bucket_function.h"
 #include "paimon/data/decimal.h"
 #include "paimon/data/timestamp.h"
 #include "paimon/memory/memory_pool.h"
@@ -236,8 +240,21 @@ static Result<WriteFunction> WriteBucketRow(int32_t col_id,
 }
 }  // namespace
 
+BucketIdCalculator::BucketIdCalculator(int32_t num_buckets,
+                                       std::unique_ptr<BucketFunction> bucket_function,
+                                       const std::shared_ptr<MemoryPool>& pool)
+    : num_buckets_(num_buckets), bucket_function_(std::move(bucket_function)), pool_(pool) {}
+
+BucketIdCalculator::~BucketIdCalculator() = default;
+
 Result<std::unique_ptr<BucketIdCalculator>> BucketIdCalculator::Create(
     bool is_pk_table, int32_t num_buckets, const std::shared_ptr<MemoryPool>& pool) {
+    return Create(is_pk_table, num_buckets, std::make_unique<DefaultBucketFunction>(), pool);
+}
+
+Result<std::unique_ptr<BucketIdCalculator>> BucketIdCalculator::Create(
+    bool is_pk_table, int32_t num_buckets, std::unique_ptr<BucketFunction> bucket_function,
+    const std::shared_ptr<MemoryPool>& pool) {
     if (num_buckets == 0 || num_buckets < -2) {
         return Status::Invalid("num buckets must be -1 or -2 or greater than 0");
     }
@@ -249,12 +266,22 @@ Result<std::unique_ptr<BucketIdCalculator>> BucketIdCalculator::Create(
     if (!is_pk_table && num_buckets == -2) {
         return Status::Invalid("Append table not support PostponeBucketMode");
     }
-    return std::unique_ptr<BucketIdCalculator>(new BucketIdCalculator(num_buckets, pool));
+    return std::unique_ptr<BucketIdCalculator>(
+        new BucketIdCalculator(num_buckets, std::move(bucket_function), pool));
+}
+
+Result<std::unique_ptr<BucketIdCalculator>> BucketIdCalculator::CreateMod(
+    bool is_pk_table, int32_t num_buckets, FieldType bucket_key_type,
+    const std::shared_ptr<MemoryPool>& pool) {
+    PAIMON_ASSIGN_OR_RAISE(auto mod_func, ModBucketFunction::Create(bucket_key_type));
+    return Create(is_pk_table, num_buckets, std::move(mod_func), pool);
 }
 
-Result<std::unique_ptr<BucketIdCalculator>> BucketIdCalculator::Create(bool is_pk_table,
-                                                                       int32_t num_buckets) {
-    return Create(is_pk_table, num_buckets, GetDefaultPool());
+Result<std::unique_ptr<BucketIdCalculator>> BucketIdCalculator::CreateHive(
+    bool is_pk_table, int32_t num_buckets, const std::vector<HiveFieldInfo>& field_infos,
+    const std::shared_ptr<MemoryPool>& pool) {
+    PAIMON_ASSIGN_OR_RAISE(auto hive_func, HiveBucketFunction::Create(field_infos));
+    return Create(is_pk_table, num_buckets, std::move(hive_func), pool);
 }
 
 Status BucketIdCalculator::CalculateBucketIds(ArrowArray* bucket_keys, ArrowSchema* bucket_schema,
@@ -298,7 +325,7 @@ Status BucketIdCalculator::CalculateBucketIds(ArrowArray* bucket_keys, ArrowSche
             write_functions[col](row, &row_writer);
         }
         row_writer.Complete();
-        bucket_ids[row] = std::abs(bucket_row.HashCode() % num_buckets_);
+        bucket_ids[row] = bucket_function_->Bucket(bucket_row, num_buckets_);
     }
     guard.Release();
     return Status::OK();