This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-rust.git
The following commit(s) were added to refs/heads/main by this push:
new d24158e feat(table): support mod and hive bucket functions (#396)
d24158e is described below
commit d24158e84ba5cb0ce7873ed769e3665061f4036f
Author: WenjunMin <[email protected]>
AuthorDate: Fri Jun 19 20:36:52 2026 +0800
feat(table): support mod and hive bucket functions (#396)
---
crates/paimon/src/spec/binary_row.rs | 2 +-
crates/paimon/src/spec/core_options.rs | 43 ++-
crates/paimon/src/table/bucket_assigner_fixed.rs | 18 +-
crates/paimon/src/table/bucket_filter.rs | 7 +-
crates/paimon/src/table/bucket_function.rs | 365 +++++++++++++++++++++++
crates/paimon/src/table/mod.rs | 1 +
crates/paimon/src/table/read_builder.rs | 4 -
crates/paimon/src/table/table_scan.rs | 131 +++++---
crates/paimon/src/table/table_write.rs | 125 +++++++-
docs/src/sql.md | 1 +
10 files changed, 635 insertions(+), 62 deletions(-)
diff --git a/crates/paimon/src/spec/binary_row.rs
b/crates/paimon/src/spec/binary_row.rs
index 104e343..145f64c 100644
--- a/crates/paimon/src/spec/binary_row.rs
+++ b/crates/paimon/src/spec/binary_row.rs
@@ -1131,7 +1131,7 @@ fn write_typed_value(
/// Build BinaryRows for all rows in the batch for the given field indices.
/// Downcasts columns once, then iterates rows — O(F) downcasts instead of
O(N*F).
-fn batch_build_binary_rows(
+pub(crate) fn batch_build_binary_rows(
batch: &RecordBatch,
field_indices: &[usize],
fields: &[crate::spec::DataField],
diff --git a/crates/paimon/src/spec/core_options.rs
b/crates/paimon/src/spec/core_options.rs
index a25fa5c..0d4e81f 100644
--- a/crates/paimon/src/spec/core_options.rs
+++ b/crates/paimon/src/spec/core_options.rs
@@ -102,6 +102,26 @@ pub enum ChangelogProducer {
Lookup,
}
+/// Bucket function used to map bucket keys to fixed bucket ids.
+///
+/// Reference: Java `CoreOptions.BucketFunctionType`.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum BucketFunctionType {
+ Default,
+ Mod,
+ Hive,
+}
+
+impl BucketFunctionType {
+ pub fn as_str(&self) -> &'static str {
+ match self {
+ Self::Default => "default",
+ Self::Mod => "mod",
+ Self::Hive => "hive",
+ }
+ }
+}
+
impl ChangelogProducer {
pub fn as_str(&self) -> &'static str {
match self {
@@ -388,15 +408,22 @@ impl<'a> CoreOptions<'a> {
.unwrap_or(DEFAULT_BUCKET)
}
- /// Whether the bucket function type is the default hash-based function.
- ///
- /// Only the default function (`Math.abs(hash % numBuckets)`) is supported
- /// for bucket predicate pruning. `mod` and `hive` use different
algorithms.
- pub fn is_default_bucket_function(&self) -> bool {
- self.options
+ /// Bucket function type. Defaults to Java-compatible Paimon hash.
+ pub fn bucket_function_type(&self) -> crate::Result<BucketFunctionType> {
+ match self
+ .options
.get(BUCKET_FUNCTION_TYPE_OPTION)
- .map(|v| v.eq_ignore_ascii_case("default"))
- .unwrap_or(true)
+ .map(|v| v.to_ascii_lowercase())
+ .as_deref()
+ .unwrap_or("default")
+ {
+ "default" => Ok(BucketFunctionType::Default),
+ "mod" => Ok(BucketFunctionType::Mod),
+ "hive" => Ok(BucketFunctionType::Hive),
+ other => Err(crate::Error::ConfigInvalid {
+ message: format!("Unsupported bucket-function.type: {other}"),
+ }),
+ }
}
/// Target file size for data files. Default is 128MB.
diff --git a/crates/paimon/src/table/bucket_assigner_fixed.rs
b/crates/paimon/src/table/bucket_assigner_fixed.rs
index a31e720..428b07d 100644
--- a/crates/paimon/src/table/bucket_assigner_fixed.rs
+++ b/crates/paimon/src/table/bucket_assigner_fixed.rs
@@ -19,9 +19,10 @@
use crate::io::FileIO;
use crate::spec::{
- batch_hash_codes, batch_to_serialized_bytes, DataField, IndexFileMeta,
EMPTY_SERIALIZED_ROW,
+ batch_to_serialized_bytes, BucketFunctionType, DataField, IndexFileMeta,
EMPTY_SERIALIZED_ROW,
};
use crate::table::bucket_assigner::{BatchAssignOutput, BucketAssigner,
PartitionBucketKey};
+use crate::table::bucket_function::batch_bucket_ids;
use crate::Result;
use arrow_array::RecordBatch;
use std::collections::HashMap;
@@ -33,6 +34,7 @@ use std::collections::HashMap;
pub(crate) struct FixedBucketAssigner {
partition_field_indices: Vec<usize>,
bucket_key_indices: Vec<usize>,
+ bucket_function_type: BucketFunctionType,
total_buckets: i32,
}
@@ -40,11 +42,13 @@ impl FixedBucketAssigner {
pub fn new(
partition_field_indices: Vec<usize>,
bucket_key_indices: Vec<usize>,
+ bucket_function_type: BucketFunctionType,
total_buckets: i32,
) -> Self {
Self {
partition_field_indices,
bucket_key_indices,
+ bucket_function_type,
total_buckets,
}
}
@@ -63,11 +67,13 @@ impl BucketAssigner for FixedBucketAssigner {
batch_to_serialized_bytes(batch, &self.partition_field_indices,
fields)?
};
- let hash_codes = batch_hash_codes(batch, &self.bucket_key_indices,
fields)?;
- let buckets: Vec<i32> = hash_codes
- .iter()
- .map(|h| (h % self.total_buckets).wrapping_abs())
- .collect();
+ let buckets = batch_bucket_ids(
+ batch,
+ &self.bucket_key_indices,
+ fields,
+ self.bucket_function_type,
+ self.total_buckets,
+ )?;
Ok(BatchAssignOutput {
partition_bytes,
diff --git a/crates/paimon/src/table/bucket_filter.rs
b/crates/paimon/src/table/bucket_filter.rs
index 7091578..238cda4 100644
--- a/crates/paimon/src/table/bucket_filter.rs
+++ b/crates/paimon/src/table/bucket_filter.rs
@@ -18,8 +18,10 @@
//! Bucket and partition predicate extraction and bucket hash pruning.
use crate::spec::{
- field_idx_to_partition_idx, BinaryRow, DataField, DataType, Datum,
Predicate, PredicateOperator,
+ field_idx_to_partition_idx, BucketFunctionType, DataField, DataType,
Datum, Predicate,
+ PredicateOperator,
};
+use crate::table::bucket_function::bucket_for_datums;
use std::collections::HashSet;
pub(super) fn split_partition_and_data_predicates(
@@ -94,6 +96,7 @@ pub(super) fn extract_predicate_for_keys(
pub(super) fn compute_target_buckets(
bucket_predicate: &Predicate,
bucket_key_fields: &[DataField],
+ bucket_function_type: BucketFunctionType,
total_buckets: i32,
) -> Option<HashSet<i32>> {
if total_buckets <= 0 || bucket_key_fields.is_empty() {
@@ -125,7 +128,7 @@ pub(super) fn compute_target_buckets(
})
.collect();
- let bucket = BinaryRow::compute_bucket_from_datums(&datums,
total_buckets);
+ let bucket = bucket_for_datums(&datums, bucket_function_type,
total_buckets).ok()?;
buckets.insert(bucket);
// Advance the combination counter (rightmost first).
diff --git a/crates/paimon/src/table/bucket_function.rs
b/crates/paimon/src/table/bucket_function.rs
new file mode 100644
index 0000000..b21082f
--- /dev/null
+++ b/crates/paimon/src/table/bucket_function.rs
@@ -0,0 +1,365 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::spec::{
+ batch_build_binary_rows, BinaryRow, BucketFunctionType, DataField,
DataType, Datum,
+};
+use arrow_array::RecordBatch;
+
+pub(crate) fn validate_bucket_function(
+ bucket_function_type: BucketFunctionType,
+ bucket_key_fields: &[DataField],
+) -> crate::Result<()> {
+ if bucket_function_type == BucketFunctionType::Mod {
+ if bucket_key_fields.len() != 1 {
+ return Err(crate::Error::ConfigInvalid {
+ message: "bucket key must have exactly one field in mod bucket
function"
+ .to_string(),
+ });
+ }
+ let data_type = bucket_key_fields[0].data_type();
+ if !matches!(data_type, DataType::Int(_) | DataType::BigInt(_)) {
+ return Err(crate::Error::ConfigInvalid {
+ message: format!(
+ "bucket key type must be INT or BIGINT in mod bucket
function, but got {data_type:?}"
+ ),
+ });
+ }
+ }
+ Ok(())
+}
+
+pub(crate) fn batch_bucket_ids(
+ batch: &RecordBatch,
+ field_indices: &[usize],
+ fields: &[DataField],
+ bucket_function_type: BucketFunctionType,
+ total_buckets: i32,
+) -> crate::Result<Vec<i32>> {
+ let rows = batch_build_binary_rows(batch, field_indices, fields)?;
+ let bucket_key_fields: Vec<DataField> = field_indices
+ .iter()
+ .map(|&idx| fields[idx].clone())
+ .collect();
+ rows.iter()
+ .map(|row| bucket_for_row(row, &bucket_key_fields,
bucket_function_type, total_buckets))
+ .collect()
+}
+
+pub(crate) fn bucket_for_datums(
+ datums: &[(Option<&Datum>, &DataType)],
+ bucket_function_type: BucketFunctionType,
+ total_buckets: i32,
+) -> crate::Result<i32> {
+ match bucket_function_type {
+ BucketFunctionType::Default => {
+ Ok(BinaryRow::compute_bucket_from_datums(datums, total_buckets))
+ }
+ BucketFunctionType::Mod => mod_bucket_from_datums(datums,
total_buckets),
+ BucketFunctionType::Hive => hive_bucket_from_datums(datums,
total_buckets),
+ }
+}
+
+fn bucket_for_row(
+ row: &BinaryRow,
+ bucket_key_fields: &[DataField],
+ bucket_function_type: BucketFunctionType,
+ total_buckets: i32,
+) -> crate::Result<i32> {
+ match bucket_function_type {
+ BucketFunctionType::Default => Ok(default_bucket(row.hash_code(),
total_buckets)),
+ BucketFunctionType::Mod => mod_bucket_from_row(row, bucket_key_fields,
total_buckets),
+ BucketFunctionType::Hive => hive_bucket_from_row(row,
bucket_key_fields, total_buckets),
+ }
+}
+
+fn default_bucket(hash: i32, total_buckets: i32) -> i32 {
+ (hash % total_buckets).wrapping_abs()
+}
+
+fn floor_mod_i64(value: i64, divisor: i32) -> i32 {
+ value.rem_euclid(divisor as i64) as i32
+}
+
+fn mod_bucket_from_row(
+ row: &BinaryRow,
+ bucket_key_fields: &[DataField],
+ total_buckets: i32,
+) -> crate::Result<i32> {
+ validate_bucket_function(BucketFunctionType::Mod, bucket_key_fields)?;
+ if row.is_null_at(0) {
+ return Ok(0);
+ }
+ match bucket_key_fields[0].data_type() {
+ DataType::Int(_) => Ok(floor_mod_i64(row.get_int(0)? as i64,
total_buckets)),
+ DataType::BigInt(_) => Ok(floor_mod_i64(row.get_long(0)?,
total_buckets)),
+ other => Err(crate::Error::Unsupported {
+ message: format!("bucket key type must be INT or BIGINT, but got
{other:?}"),
+ }),
+ }
+}
+
+fn mod_bucket_from_datums(
+ datums: &[(Option<&Datum>, &DataType)],
+ total_buckets: i32,
+) -> crate::Result<i32> {
+ if datums.len() != 1 {
+ return Err(crate::Error::ConfigInvalid {
+ message: "bucket key must have exactly one field in mod bucket
function".to_string(),
+ });
+ }
+ match datums[0] {
+ (None, DataType::Int(_) | DataType::BigInt(_)) => Ok(0),
+ (Some(Datum::Int(v)), DataType::Int(_)) => Ok(floor_mod_i64(*v as i64,
total_buckets)),
+ (Some(Datum::Long(v)), DataType::BigInt(_)) => Ok(floor_mod_i64(*v,
total_buckets)),
+ (_, data_type) => Err(crate::Error::Unsupported {
+ message: format!("bucket key type must be INT or BIGINT, but got
{data_type:?}"),
+ }),
+ }
+}
+
+fn hive_bucket_from_row(
+ row: &BinaryRow,
+ bucket_key_fields: &[DataField],
+ total_buckets: i32,
+) -> crate::Result<i32> {
+ let mut hash = 0_i32;
+ for (pos, field) in bucket_key_fields.iter().enumerate() {
+ let datum = row.get_datum(pos, field.data_type())?;
+ hash = hash
+ .wrapping_mul(31)
+ .wrapping_add(hive_hash_datum(datum.as_ref(), field.data_type())?);
+ }
+ Ok(positive_mod(hash, total_buckets))
+}
+
+fn hive_bucket_from_datums(
+ datums: &[(Option<&Datum>, &DataType)],
+ total_buckets: i32,
+) -> crate::Result<i32> {
+ let mut hash = 0_i32;
+ for (datum, data_type) in datums {
+ hash = hash
+ .wrapping_mul(31)
+ .wrapping_add(hive_hash_datum(*datum, data_type)?);
+ }
+ Ok(positive_mod(hash, total_buckets))
+}
+
+fn positive_mod(hash: i32, total_buckets: i32) -> i32 {
+ ((hash as u32 & 0x7fff_ffff) % total_buckets as u32) as i32
+}
+
+fn hive_hash_datum(datum: Option<&Datum>, data_type: &DataType) ->
crate::Result<i32> {
+ let Some(datum) = datum else {
+ return Ok(0);
+ };
+
+ match (datum, data_type) {
+ (Datum::Bool(v), DataType::Boolean(_)) => Ok(i32::from(*v)),
+ (Datum::TinyInt(v), DataType::TinyInt(_)) => Ok(*v as i32),
+ (Datum::SmallInt(v), DataType::SmallInt(_)) => Ok(*v as i32),
+ (Datum::Int(v), DataType::Int(_)) => Ok(*v),
+ (Datum::Long(v), DataType::BigInt(_)) => Ok(java_long_hash(*v)),
+ (Datum::Float(v), DataType::Float(_)) => Ok(java_float_bits(*v) as
i32),
+ (Datum::Double(v), DataType::Double(_)) =>
Ok(java_long_hash(java_double_bits(*v) as i64)),
+ (Datum::String(v), DataType::Char(_) | DataType::VarChar(_)) => {
+ Ok(hive_hash_bytes(v.as_bytes()))
+ }
+ (Datum::Bytes(v), DataType::Binary(_) | DataType::VarBinary(_)) =>
Ok(hive_hash_bytes(v)),
+ (
+ Datum::Decimal {
+ unscaled, scale, ..
+ },
+ DataType::Decimal(_),
+ ) => {
+ let (unscaled, scale) = normalize_decimal(*unscaled, *scale);
+ Ok(java_big_decimal_hash(unscaled, scale))
+ }
+ (Datum::Date(v), DataType::Date(_)) => Ok(*v),
+ (Datum::Time(v), DataType::Time(_)) => Ok(*v),
+ _ => Err(crate::Error::Unsupported {
+ message: format!("Unsupported type as bucket key type
{data_type:?}"),
+ }),
+ }
+}
+
+fn hive_hash_bytes(bytes: &[u8]) -> i32 {
+ bytes.iter().fold(0_i32, |hash, byte| {
+ hash.wrapping_mul(31).wrapping_add(*byte as i8 as i32)
+ })
+}
+
+fn java_long_hash(value: i64) -> i32 {
+ let bits = value as u64;
+ (bits ^ (bits >> 32)) as u32 as i32
+}
+
+fn java_float_bits(value: f32) -> u32 {
+ if value == 0.0 {
+ 0
+ } else if value.is_nan() {
+ 0x7fc0_0000
+ } else {
+ value.to_bits()
+ }
+}
+
+fn java_double_bits(value: f64) -> u64 {
+ if value == 0.0 {
+ 0
+ } else if value.is_nan() {
+ 0x7ff8_0000_0000_0000
+ } else {
+ value.to_bits()
+ }
+}
+
+fn normalize_decimal(mut unscaled: i128, mut scale: u32) -> (i128, u32) {
+ if unscaled == 0 {
+ return (0, 0);
+ }
+ while scale > 0 && unscaled % 10 == 0 {
+ unscaled /= 10;
+ scale -= 1;
+ }
+ (unscaled, scale)
+}
+
+fn java_big_decimal_hash(unscaled: i128, scale: u32) -> i32 {
+ if let Ok(compact) = i64::try_from(unscaled) {
+ if compact != i64::MIN {
+ let val = if compact < 0 {
+ compact.wrapping_neg() as u64
+ } else {
+ compact as u64
+ };
+ let temp = ((val >> 32) as i32)
+ .wrapping_mul(31)
+ .wrapping_add(val as u32 as i32);
+ let signed_temp = if compact < 0 {
+ temp.wrapping_neg()
+ } else {
+ temp
+ };
+ return signed_temp.wrapping_mul(31).wrapping_add(scale as i32);
+ }
+ }
+
+ java_big_integer_hash(unscaled)
+ .wrapping_mul(31)
+ .wrapping_add(scale as i32)
+}
+
+fn java_big_integer_hash(value: i128) -> i32 {
+ if value == 0 {
+ return 0;
+ }
+
+ let sign = if value < 0 { -1_i32 } else { 1_i32 };
+ let mut magnitude = value.unsigned_abs();
+ let mut words = Vec::new();
+ while magnitude != 0 {
+ words.push((magnitude & 0xffff_ffff) as u32);
+ magnitude >>= 32;
+ }
+ words.reverse();
+
+ let hash = words.into_iter().fold(0_i32, |hash, word| {
+ hash.wrapping_mul(31).wrapping_add(word as i32)
+ });
+ hash.wrapping_mul(sign)
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::spec::{BigIntType, BooleanType, DecimalType, IntType,
VarBinaryType, VarCharType};
+
+ #[test]
+ fn mod_bucket_uses_floor_mod_for_int_and_bigint() {
+ let int_type = DataType::Int(IntType::new());
+ let long_type = DataType::BigInt(BigIntType::new());
+
+ assert_eq!(
+ bucket_for_datums(
+ &[(Some(&Datum::Int(-3)), &int_type)],
+ BucketFunctionType::Mod,
+ 5,
+ )
+ .unwrap(),
+ 2
+ );
+ assert_eq!(
+ bucket_for_datums(
+ &[(Some(&Datum::Long(17)), &long_type)],
+ BucketFunctionType::Mod,
+ 5,
+ )
+ .unwrap(),
+ 2
+ );
+ }
+
+ #[test]
+ fn hive_bucket_matches_java_reference_case() {
+ let bool_type = DataType::Boolean(BooleanType::new());
+ let int_type = DataType::Int(IntType::new());
+ let string_type = DataType::VarChar(VarCharType::default());
+ let bytes_type = DataType::VarBinary(VarBinaryType::default());
+ let decimal_type = DataType::Decimal(DecimalType::new(10, 4).unwrap());
+
+ let bucket = bucket_for_datums(
+ &[
+ (Some(&Datum::Bool(true)), &bool_type),
+ (Some(&Datum::Int(7)), &int_type),
+ (Some(&Datum::String("hello".into())), &string_type),
+ (Some(&Datum::Bytes(vec![1, 2, 3])), &bytes_type),
+ (
+ Some(&Datum::Decimal {
+ unscaled: 123400,
+ precision: 10,
+ scale: 4,
+ }),
+ &decimal_type,
+ ),
+ ],
+ BucketFunctionType::Hive,
+ 8,
+ )
+ .unwrap();
+
+ let expected_hash = 31_i32
+ .wrapping_mul(
+ 31_i32
+ .wrapping_mul(
+ 31_i32
+
.wrapping_mul(31_i32.wrapping_mul(1).wrapping_add(7))
+ .wrapping_add(99_162_322),
+ )
+ .wrapping_add(1_026),
+ )
+ .wrapping_add(38_256);
+ assert_eq!(bucket, positive_mod(expected_hash, 8));
+ }
+
+ #[test]
+ fn hive_decimal_hash_trims_trailing_zeros_like_big_decimal() {
+ assert_eq!(normalize_decimal(123400, 4), (1234, 2));
+ assert_eq!(normalize_decimal(0, 8), (0, 0));
+ }
+}
diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs
index 3158ffc..e7a9052 100644
--- a/crates/paimon/src/table/mod.rs
+++ b/crates/paimon/src/table/mod.rs
@@ -27,6 +27,7 @@ mod bucket_assigner_cross;
mod bucket_assigner_dynamic;
mod bucket_assigner_fixed;
mod bucket_filter;
+mod bucket_function;
mod commit_message;
pub(crate) mod cow_writer;
mod data_evolution_reader;
diff --git a/crates/paimon/src/table/read_builder.rs
b/crates/paimon/src/table/read_builder.rs
index f11acc7..099a8a0 100644
--- a/crates/paimon/src/table/read_builder.rs
+++ b/crates/paimon/src/table/read_builder.rs
@@ -70,10 +70,6 @@ pub(super) fn split_scan_predicates(
fn bucket_predicate(table: &Table, filter: &Predicate) -> Option<Predicate> {
let core_options = CoreOptions::new(table.schema().options());
- if !core_options.is_default_bucket_function() {
- return None;
- }
-
let bucket_keys = core_options.bucket_key().unwrap_or_else(|| {
if table.schema().trimmed_primary_keys().is_empty() {
Vec::new()
diff --git a/crates/paimon/src/table/table_scan.rs
b/crates/paimon/src/table/table_scan.rs
index 170405b..bd847d1 100644
--- a/crates/paimon/src/table/table_scan.rs
+++ b/crates/paimon/src/table/table_scan.rs
@@ -30,8 +30,9 @@ use super::stats_filter::{
use super::Table;
use crate::io::FileIO;
use crate::spec::{
- avro::SharedSchemaCache, bucket_dir_name, BinaryRow, CoreOptions,
DataField, DataFileMeta,
- FileKind, IndexManifest, ManifestEntry, PartitionComputer, Predicate,
Snapshot,
+ avro::SharedSchemaCache, bucket_dir_name, BinaryRow, BucketFunctionType,
CoreOptions,
+ DataField, DataFileMeta, FileKind, IndexManifest, ManifestEntry,
PartitionComputer, Predicate,
+ Snapshot,
};
use crate::table::bin_pack::split_for_batch;
use crate::table::merge_tree_split_generator::{
@@ -92,6 +93,7 @@ async fn read_all_manifest_entries(
schema_fields: &[DataField],
bucket_predicate: Option<&Predicate>,
bucket_key_fields: &[DataField],
+ bucket_function_type: BucketFunctionType,
) -> crate::Result<Vec<ManifestEntry>> {
let (mut manifest_files, delta) = futures::try_join!(
read_manifest_list(file_io, table_path, snapshot.base_manifest_list()),
@@ -142,7 +144,12 @@ async fn read_all_manifest_entries(
}
if let Some(pred) = bucket_predicate {
let targets =
bucket_cache.entry(total_buckets).or_insert_with(|| {
- compute_target_buckets(pred,
bucket_key_fields, total_buckets)
+ compute_target_buckets(
+ pred,
+ bucket_key_fields,
+ bucket_function_type,
+ total_buckets,
+ )
});
if let Some(targets) = targets {
if !targets.contains(&bucket) {
@@ -489,29 +496,29 @@ impl<'a> TableScan<'a> {
self.data_predicates.as_slice()
};
- let bucket_key_fields: Vec<DataField> =
- if self.bucket_predicate.is_none() ||
!core_options.is_default_bucket_function() {
- Vec::new()
- } else {
- let bucket_keys = core_options.bucket_key().unwrap_or_else(|| {
- if has_primary_keys {
- self.table.schema().trimmed_primary_keys()
- } else {
- Vec::new()
- }
- });
- bucket_keys
- .iter()
- .filter_map(|key| {
- self.table
- .schema()
- .fields()
- .iter()
- .find(|f| f.name() == key)
- .cloned()
- })
- .collect::<Vec<_>>()
- };
+ let bucket_key_fields: Vec<DataField> = if
self.bucket_predicate.is_none() {
+ Vec::new()
+ } else {
+ let bucket_keys = core_options.bucket_key().unwrap_or_else(|| {
+ if has_primary_keys {
+ self.table.schema().trimmed_primary_keys()
+ } else {
+ Vec::new()
+ }
+ });
+ bucket_keys
+ .iter()
+ .filter_map(|key| {
+ self.table
+ .schema()
+ .fields()
+ .iter()
+ .find(|f| f.name() == key)
+ .cloned()
+ })
+ .collect::<Vec<_>>()
+ };
+ let bucket_function_type = core_options.bucket_function_type()?;
let entries = read_all_manifest_entries(
file_io,
@@ -527,6 +534,7 @@ impl<'a> TableScan<'a> {
self.table.schema().fields(),
self.bucket_predicate.as_ref(),
&bucket_key_fields,
+ bucket_function_type,
)
.await?;
Ok(merge_manifest_entries(entries))
@@ -827,10 +835,10 @@ mod tests {
use crate::catalog::Identifier;
use crate::io::FileIOBuilder;
use crate::spec::{
- stats::BinaryTableStats, ArrayType, BinaryRow, BinaryRowBuilder,
DataField, DataFileMeta,
- DataType, Datum, DeletionVectorMeta, FileKind, IndexFileMeta,
IndexManifestEntry, IntType,
- Predicate, PredicateBuilder, PredicateOperator, Schema as
PaimonSchema, TableSchema,
- VarCharType,
+ stats::BinaryTableStats, ArrayType, BinaryRow, BinaryRowBuilder,
BucketFunctionType,
+ DataField, DataFileMeta, DataType, Datum, DeletionVectorMeta,
FileKind, IndexFileMeta,
+ IndexManifestEntry, IntType, Predicate, PredicateBuilder,
PredicateOperator,
+ Schema as PaimonSchema, TableSchema, VarCharType,
};
use crate::table::bucket_filter::{compute_target_buckets,
extract_predicate_for_keys};
use crate::table::partition_filter::PartitionFilter;
@@ -839,6 +847,7 @@ mod tests {
use crate::table::Table;
use crate::Error;
use chrono::{DateTime, Utc};
+ use std::collections::HashSet;
/// Helper to build a DataFileMeta with data evolution fields.
fn make_evo_file(
@@ -1524,7 +1533,7 @@ mod tests {
literals: vec![Datum::Int(42)],
};
- let buckets = compute_target_buckets(&pred, &fields, 4);
+ let buckets = compute_target_buckets(&pred, &fields,
BucketFunctionType::Default, 4);
assert!(buckets.is_some());
let buckets = buckets.unwrap();
assert_eq!(buckets.len(), 1);
@@ -1544,7 +1553,7 @@ mod tests {
literals: vec![Datum::Int(1), Datum::Int(2), Datum::Int(3)],
};
- let buckets = compute_target_buckets(&pred, &fields, 4);
+ let buckets = compute_target_buckets(&pred, &fields,
BucketFunctionType::Default, 4);
assert!(buckets.is_some());
let buckets = buckets.unwrap();
// Should have at most 3 buckets (could be fewer if some hash to the
same bucket)
@@ -1566,7 +1575,7 @@ mod tests {
literals: vec![Datum::Int(10)],
};
- let buckets = compute_target_buckets(&pred, &fields, 4);
+ let buckets = compute_target_buckets(&pred, &fields,
BucketFunctionType::Default, 4);
assert!(
buckets.is_none(),
"Range predicates cannot determine target buckets"
@@ -1596,7 +1605,7 @@ mod tests {
},
]);
- let buckets = compute_target_buckets(&pred, &fields, 8);
+ let buckets = compute_target_buckets(&pred, &fields,
BucketFunctionType::Default, 8);
assert!(buckets.is_some());
let buckets = buckets.unwrap();
assert_eq!(buckets.len(), 1);
@@ -1619,7 +1628,7 @@ mod tests {
literals: vec![Datum::Int(1)],
};
- let buckets = compute_target_buckets(&pred, &fields, 8);
+ let buckets = compute_target_buckets(&pred, &fields,
BucketFunctionType::Default, 8);
assert!(
buckets.is_none(),
"Partial bucket key should not determine target buckets"
@@ -1641,7 +1650,7 @@ mod tests {
literals: vec![Datum::String("alice".into())],
};
- let buckets = compute_target_buckets(&pred, &fields, 4);
+ let buckets = compute_target_buckets(&pred, &fields,
BucketFunctionType::Default, 4);
assert!(buckets.is_some());
let buckets = buckets.unwrap();
assert_eq!(buckets.len(), 1);
@@ -1649,6 +1658,52 @@ mod tests {
assert!((0..4).contains(&bucket));
}
+ #[test]
+ fn test_compute_target_buckets_mod_function() {
+ let fields = bucket_key_fields();
+ let pred = Predicate::Leaf {
+ column: "id".into(),
+ index: 0,
+ data_type: DataType::Int(IntType::new()),
+ op: PredicateOperator::Eq,
+ literals: vec![Datum::Int(-3)],
+ };
+
+ let buckets = compute_target_buckets(&pred, &fields,
BucketFunctionType::Mod, 5);
+ assert_eq!(buckets, Some(HashSet::from([2])));
+ }
+
+ #[test]
+ fn test_compute_target_buckets_hive_function() {
+ let fields = vec![
+ DataField::new(0, "id".to_string(), DataType::Int(IntType::new())),
+ DataField::new(
+ 1,
+ "name".to_string(),
+ DataType::VarChar(VarCharType::default()),
+ ),
+ ];
+ let pred = Predicate::And(vec![
+ Predicate::Leaf {
+ column: "id".into(),
+ index: 0,
+ data_type: DataType::Int(IntType::new()),
+ op: PredicateOperator::Eq,
+ literals: vec![Datum::Int(7)],
+ },
+ Predicate::Leaf {
+ column: "name".into(),
+ index: 1,
+ data_type: DataType::VarChar(VarCharType::default()),
+ op: PredicateOperator::Eq,
+ literals: vec![Datum::String("hello".into())],
+ },
+ ]);
+
+ let buckets = compute_target_buckets(&pred, &fields,
BucketFunctionType::Hive, 8);
+ assert_eq!(buckets, Some(HashSet::from([3])));
+ }
+
#[test]
fn test_compute_target_buckets_is_null() {
let fields = bucket_key_fields();
@@ -1660,7 +1715,7 @@ mod tests {
literals: vec![],
};
- let buckets = compute_target_buckets(&pred, &fields, 4);
+ let buckets = compute_target_buckets(&pred, &fields,
BucketFunctionType::Default, 4);
assert!(buckets.is_some(), "IsNull should determine a target bucket");
let buckets = buckets.unwrap();
assert_eq!(buckets.len(), 1);
@@ -1698,7 +1753,7 @@ mod tests {
},
]);
- let buckets = compute_target_buckets(&pred, &fields, 8);
+ let buckets = compute_target_buckets(&pred, &fields,
BucketFunctionType::Default, 8);
assert!(
buckets.is_some(),
"Composite key with IsNull should determine a target bucket"
diff --git a/crates/paimon/src/table/table_write.rs
b/crates/paimon/src/table/table_write.rs
index e5b62e4..2eaa6d2 100644
--- a/crates/paimon/src/table/table_write.rs
+++ b/crates/paimon/src/table/table_write.rs
@@ -23,8 +23,8 @@
use crate::arrow::build_target_arrow_schema;
use crate::spec::PartitionComputer;
use crate::spec::{
- first_row_supports_changelog_producer, BinaryRow, ChangelogProducer,
CoreOptions, DataType,
- MergeEngine, EMPTY_SERIALIZED_ROW, POSTPONE_BUCKET,
+ first_row_supports_changelog_producer, BinaryRow, ChangelogProducer,
CoreOptions, DataField,
+ DataType, MergeEngine, EMPTY_SERIALIZED_ROW, POSTPONE_BUCKET,
};
use crate::table::blob_file_writer::AppendBlobFileWriter;
use crate::table::bucket_assigner::{BucketAssignerEnum, PartitionBucketKey};
@@ -32,6 +32,7 @@ use
crate::table::bucket_assigner_constant::ConstantBucketAssigner;
use crate::table::bucket_assigner_cross::CrossPartitionAssigner;
use crate::table::bucket_assigner_dynamic::DynamicBucketAssigner;
use crate::table::bucket_assigner_fixed::FixedBucketAssigner;
+use crate::table::bucket_function::validate_bucket_function;
use crate::table::commit_message::CommitMessage;
use crate::table::data_file_writer::DataFileWriter;
use crate::table::kv_file_writer::{KeyValueFileWriter, KeyValueWriteConfig};
@@ -259,6 +260,7 @@ impl TableWrite {
}
let target_bucket_row_number =
core_options.dynamic_bucket_target_row_num();
+ let bucket_function_type = core_options.bucket_function_type()?;
let bucket_assigner = if is_dynamic_cross_partition {
BucketAssignerEnum::CrossPartition(Box::new(CrossPartitionAssigner::new(
@@ -286,9 +288,15 @@ impl TableWrite {
} else if total_buckets <= 1 || bucket_key_indices.is_empty() {
BucketAssignerEnum::Constant(ConstantBucketAssigner::new(partition_field_indices,
0))
} else {
+ let bucket_key_fields: Vec<DataField> = bucket_key_indices
+ .iter()
+ .map(|&idx| fields[idx].clone())
+ .collect();
+ validate_bucket_function(bucket_function_type,
&bucket_key_fields)?;
BucketAssignerEnum::Fixed(FixedBucketAssigner::new(
partition_field_indices,
bucket_key_indices,
+ bucket_function_type,
total_buckets,
))
};
@@ -740,7 +748,7 @@ mod tests {
};
use crate::table::{SnapshotManager, TableCommit};
use arrow_array::RecordBatchReader as _;
- use arrow_array::{Int32Array, Int64Array, Int8Array};
+ use arrow_array::{Int32Array, Int64Array, Int8Array, StringArray};
use arrow_schema::{
DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema,
TimeUnit,
};
@@ -1319,6 +1327,117 @@ mod tests {
.unwrap()
}
+ #[tokio::test]
+ async fn test_mod_bucket_function_routes_by_floor_mod() {
+ let file_io = test_file_io();
+ let table_path = "memory:/test_mod_bucket_function";
+ setup_dirs(&file_io, table_path).await;
+
+ let schema = Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .column("value", DataType::Int(IntType::new()))
+ .option("bucket", "5")
+ .option("bucket-key", "id")
+ .option("bucket-function.type", "mod")
+ .build()
+ .unwrap();
+ let table = Table::new(
+ file_io.clone(),
+ Identifier::new("default", "test_table"),
+ table_path.to_string(),
+ TableSchema::new(0, &schema),
+ None,
+ );
+
+ let fields = table.schema().fields().to_vec();
+ let mut table_write = TableWrite::new(&table,
"test-user".to_string()).unwrap();
+ let output = table_write
+ .bucket_assigner
+ .assign_batch(&make_batch(vec![-3, 17, 5], vec![10, 20, 30]),
&fields)
+ .await
+ .unwrap();
+
+ assert_eq!(output.buckets, vec![2, 2, 0]);
+ }
+
+ #[tokio::test]
+ async fn test_mod_bucket_function_rejects_non_integral_key() {
+ let file_io = test_file_io();
+ let table_path = "memory:/test_mod_bucket_invalid";
+ setup_dirs(&file_io, table_path).await;
+
+ let schema = Schema::builder()
+ .column("id", DataType::VarChar(VarCharType::string_type()))
+ .column("value", DataType::Int(IntType::new()))
+ .option("bucket", "5")
+ .option("bucket-key", "id")
+ .option("bucket-function.type", "mod")
+ .build()
+ .unwrap();
+ let table = Table::new(
+ file_io,
+ Identifier::new("default", "test_table"),
+ table_path.to_string(),
+ TableSchema::new(0, &schema),
+ None,
+ );
+
+ let err = match TableWrite::new(&table, "test-user".to_string()) {
+ Ok(_) => panic!("expected mod bucket function to reject
non-integral key"),
+ Err(err) => err,
+ };
+ assert!(
+ matches!(err, crate::Error::ConfigInvalid { ref message } if
message.contains("INT or BIGINT")),
+ "expected ConfigInvalid for non-integral mod bucket key, got
{err:?}"
+ );
+ }
+
+ #[tokio::test]
+ async fn test_hive_bucket_function_routes_by_hive_hash() {
+ let file_io = test_file_io();
+ let table_path = "memory:/test_hive_bucket_function";
+ setup_dirs(&file_io, table_path).await;
+
+ let schema = Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .column("name", DataType::VarChar(VarCharType::string_type()))
+ .option("bucket", "8")
+ .option("bucket-key", "id,name")
+ .option("bucket-function.type", "hive")
+ .build()
+ .unwrap();
+ let table = Table::new(
+ file_io.clone(),
+ Identifier::new("default", "test_table"),
+ table_path.to_string(),
+ TableSchema::new(0, &schema),
+ None,
+ );
+
+ let arrow_schema = Arc::new(ArrowSchema::new(vec![
+ ArrowField::new("id", ArrowDataType::Int32, false),
+ ArrowField::new("name", ArrowDataType::Utf8, false),
+ ]));
+ let batch = RecordBatch::try_new(
+ arrow_schema,
+ vec![
+ Arc::new(Int32Array::from(vec![7])),
+ Arc::new(StringArray::from(vec!["hello"])),
+ ],
+ )
+ .unwrap();
+
+ let fields = table.schema().fields().to_vec();
+ let mut table_write = TableWrite::new(&table,
"test-user".to_string()).unwrap();
+ let output = table_write
+ .bucket_assigner
+ .assign_batch(&batch, &fields)
+ .await
+ .unwrap();
+
+ assert_eq!(output.buckets, vec![3]);
+ }
+
#[tokio::test]
async fn test_write_bucketed_with_null_bucket_key() {
let file_io = test_file_io();
diff --git a/docs/src/sql.md b/docs/src/sql.md
index 6705578..a4dc0fa 100644
--- a/docs/src/sql.md
+++ b/docs/src/sql.md
@@ -903,6 +903,7 @@ Set via `WITH ('key' = 'value')` at table creation time, or
dynamically via `SET
| `'bucket' = '-1'` | Dynamic bucket mode (HASH index) |
| `'bucket' = '-2'` | Postpone bucket mode (deferred assignment) |
| `'bucket-key' = 'col'` | Explicit bucket key column |
+| `'bucket-function.type' = 'default' \| 'mod' \| 'hive'` | Function used to
map fixed bucket keys to bucket ids |
### Merge Engine
