feat: implement bucket transform

For type:
- Int32
- Int64
- Date32
- Time32
- Utf8/String

Author: natir

Cargo.toml

@@ -38,6 +38,7 @@ itertools = "0.14.0"
 lazy_static = "1.5.0"
 lru = "0.16.0"
 object_store = { version = "0.12", features = ["aws", "gcp"] }
+fasthash = { version = "0.4" }
 parquet = { version = "56", features = ["async", "object_store"] }
 pin-project-lite = "0.2"
 regex = "1.11.1"

iceberg-rust/Cargo.toml

@@ -24,6 +24,7 @@ itertools = { workspace = true }
 lazy_static = { workspace = true }
 lru = { workspace = true }
 object_store = { workspace = true }
+fasthash = { workspace = true }
 parquet = { workspace = true }
 pin-project-lite = { workspace = true }
 regex = { workspace = true }
@@ -44,4 +45,3 @@ uuid = { workspace = true }
 rstest = "0.23.0"
 chrono = { workspace = true }
 iceberg-sql-catalog = { path = "../catalogs/iceberg-sql-catalog" }
-

iceberg-rust/src/arrow/transform.rs

@@ -14,7 +14,8 @@
 use std::sync::Arc;
 
 use arrow::{
-    array::{as_primitive_array, Array, ArrayRef, PrimitiveArray},
+    array::{as_primitive_array, downcast_array, Array, ArrayRef, PrimitiveArray, StringArray},
+    buffer::ScalarBuffer,
     compute::{binary, cast, date_part, unary, DatePart},
     datatypes::{
         DataType, Date32Type, Int16Type, Int32Type, Int64Type, TimeUnit, TimestampMicrosecondType,
@@ -43,6 +44,14 @@ static MICROS_IN_DAY: i64 = 86_400_000_000;
 /// * Month - Extracts month from date32 or timestamp
 /// * Year - Extracts year from date32 or timestamp
 /// * Hour - Extracts hour from timestamp
+/// * Int16 - Truncate value
+/// * Int32 - Truncate value
+/// * Int64 - Truncate value
+/// * Int32 - Use hash of value to repart it between bucket
+/// * Int64 - Use hash of value to repart it between bucket
+/// * Date32 - Use hash of value to repart it between bucket
+/// * Time32 - Use hash of value to repart it between bucket
+/// * Utf8 - Use hash of value to repart it between bucket
 pub fn transform_arrow(array: ArrayRef, transform: &Transform) -> Result<ArrayRef, ArrowError> {
     match (array.data_type(), transform) {
         (_, Transform::Identity) => Ok(array),
@@ -114,6 +123,56 @@ pub fn transform_arrow(array: ArrayRef, transform: &Transform) -> Result<ArrayRe
                 i - i.rem_euclid(*m as i64)
             }),
         )),
+        (DataType::Int32, Transform::Bucket(m)) => Ok(Arc::<PrimitiveArray<Int32Type>>::new(
+            unary(as_primitive_array::<Int32Type>(&array), |i| {
+                (fasthash::murmur3::hash32_with_seed((i as i64).to_le_bytes(), 0) as i32)
+                    .rem_euclid(*m as i32)
+            }),
+        )),
+        (DataType::Int64, Transform::Bucket(m)) => Ok(Arc::<PrimitiveArray<Int32Type>>::new(
+            unary(as_primitive_array::<Int64Type>(&array), |i| {
+                (fasthash::murmur3::hash32_with_seed(i.to_le_bytes(), 0) as i32)
+                    .rem_euclid(*m as i32)
+            }),
+        )),
+        (DataType::Date32, Transform::Bucket(m)) => {
+            let temp = cast(&array, &DataType::Int32)?;
+
+            Ok(Arc::<PrimitiveArray<Int32Type>>::new(unary(
+                as_primitive_array::<Int32Type>(&temp),
+                |i| {
+                    (fasthash::murmur3::hash32_with_seed(i.to_le_bytes(), 0) as i32)
+                        .rem_euclid(*m as i32)
+                },
+            )))
+        }
+        (DataType::Time32(TimeUnit::Millisecond), Transform::Bucket(m)) => {
+            let temp = cast(&array, &DataType::Int32)?;
+
+            Ok(Arc::<PrimitiveArray<Int32Type>>::new(unary(
+                as_primitive_array::<Int32Type>(&temp),
+                |i: i32| {
+                    (fasthash::murmur3::hash32_with_seed((i as i64).to_le_bytes(), 0) as i32)
+                        .rem_euclid(*m as i32)
+                },
+            )))
+        }
+        (DataType::Utf8, Transform::Bucket(m)) => {
+            let nulls = array.nulls();
+            let local_array: StringArray = downcast_array::<StringArray>(&array);
+
+            Ok(Arc::new(PrimitiveArray::<Int32Type>::new(
+                ScalarBuffer::from_iter(local_array.iter().map(|a| {
+                    if let Some(value) = a {
+                        fasthash::murmur3::hash32_with_seed(value.as_bytes(), 0) as i32
+                    } else {
+                        0
+                    }
+                    .rem_euclid(*m as i32)
+                })),
+                nulls.cloned(),
+            )))
+        }
         _ => Err(ArrowError::ComputeError(
             "Failed to perform transform for datatype".to_string(),
         )),
@@ -322,6 +381,123 @@ mod tests {
         assert_eq!(&expected, &result);
     }
 
+    #[test]
+    fn test_bucket_hash_value() {
+        // Check value match https://iceberg.apache.org/spec/#appendix-b-32-bit-hash-requirements
+
+        // 34 -> 2017239379
+        assert_eq!(
+            fasthash::murmur3::hash32_with_seed((34i32 as i64).to_le_bytes(), 0),
+            2017239379
+        );
+        // 34 -> 2017239379
+        assert_eq!(
+            fasthash::murmur3::hash32_with_seed((34i64).to_le_bytes(), 0),
+            2017239379
+        );
+        // daysFromUnixEpoch(2017-11-16) -> 17_486 -> -653330422
+        assert_eq!(
+            fasthash::murmur3::hash32_with_seed((17_486i32 as i64).to_le_bytes(), 0) as i32,
+            -653330422
+        );
+        // 81_068_000_000 number of micros from midnight 22:31:08
+        assert_eq!(
+            fasthash::murmur3::hash32_with_seed((81_068_000_000i64).to_le_bytes(), 0) as i32,
+            -662762989
+        );
+
+        // utf8Bytes(iceberg) -> 1210000089
+        assert_eq!(
+            fasthash::murmur3::hash32_with_seed("iceberg".as_bytes(), 0) as i32,
+            1210000089
+        );
+    }
+
+    #[test]
+    fn test_int32_bucket_transform() {
+        assert_eq!(
+            fasthash::murmur3::hash32_with_seed(17_486i64.to_le_bytes(), 0) as i32,
+            -653_330_422
+        );
+
+        let array = Arc::new(arrow::array::Int32Array::from(vec![
+            Some(34),       // Spec value
+            Some(17_486),   // number of day between 2017-11-16 and epoch
+            Some(84668000), // number of micros from midnight 22:31:08
+            Some(-2000),
+            Some(0),
+            None,
+        ])) as ArrayRef;
+        let result = transform_arrow(array, &Transform::Bucket(2_017_239_380)).unwrap();
+        let expected = Arc::new(arrow::array::Int32Array::from(vec![
+            Some(2017239379),
+            Some(1363908958),
+            Some(988822981),
+            Some(964620854),
+            Some(1669671676),
+            None,
+        ])) as ArrayRef;
+        assert_eq!(&expected, &result);
+    }
+
+    #[test]
+    fn test_int64_bucket_transform() {
+        let array = Arc::new(arrow::array::Int64Array::from(vec![
+            Some(34),     // Spec value
+            Some(17_486), // number of day between 2017-11-16 and epoch
+            Some(2000),
+            Some(-2000),
+            Some(0),
+            None,
+        ])) as ArrayRef;
+        let result = transform_arrow(array, &Transform::Bucket(2_017_239_380)).unwrap();
+        let expected = Arc::new(arrow::array::Int32Array::from(vec![
+            Some(2017239379),
+            Some(1363908958),
+            Some(716_914_497),
+            Some(964_620_854),
+            Some(1669671676),
+            None,
+        ])) as ArrayRef;
+        assert_eq!(&expected, &result);
+    }
+
+    #[test]
+    fn test_date32_bucket_transform() {
+        let array = Arc::new(arrow::array::Date32Array::from(vec![
+            Some(17_486), // number of day between 2017-11-16
+            None,
+        ])) as ArrayRef;
+        let result = transform_arrow(array, &Transform::Bucket(2_017_239_380)).unwrap();
+        let expected =
+            Arc::new(arrow::array::Int32Array::from(vec![Some(1087118125), None])) as ArrayRef;
+        assert_eq!(&expected, &result);
+    }
+
+    #[test]
+    fn test_time32_bucket_transform() {
+        let array = Arc::new(arrow::array::Time32MillisecondArray::from(vec![
+            Some(81_068_000), // number of number of micros from midnight 22:31:08
+            None,
+        ])) as ArrayRef;
+        let result = transform_arrow(array, &Transform::Bucket(29_824_073)).unwrap();
+        let expected =
+            Arc::new(arrow::array::Int32Array::from(vec![Some(10797584), None])) as ArrayRef;
+        assert_eq!(&expected, &result);
+    }
+
+    #[test]
+    fn test_utf8_bucket_transform() {
+        let array =
+            Arc::new(arrow::array::StringArray::from(vec![Some("iceberg"), None])) as ArrayRef;
+        let result = transform_arrow(array, &Transform::Bucket(1_210_000_090)).unwrap();
+        let expected = Arc::new(arrow::array::Int32Array::from(vec![
+            Some(1_210_000_089),
+            None,
+        ])) as ArrayRef;
+        assert_eq!(&expected, &result);
+    }
+
     #[test]
     fn test_unsupported_transform() {
         let array = Arc::new(arrow::array::StringArray::from(vec!["a", "b", "c"])) as ArrayRef;