Merge pull request JanKaul#263 from pierre-marijon/add_transforms

Add some bucket transformation

Author: JanKaul

Cargo.lock

@@ -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"] }
+murmur3 = { version = "0.5.2" }
 parquet = { version = "56", features = ["async", "object_store"] }
 pin-project-lite = "0.2"
 regex = "1.11.1"

Cargo.toml

@@ -16,7 +16,7 @@ derive-getters = { workspace = true }
 derive_builder = { workspace = true }
 getrandom = { workspace = true }
 itertools = { workspace = true }
-murmur3 = "0.5.2"
+murmur3 = { workspace = true }
 ordered-float = { version = "5.0.0", features = ["serde"] }
 rust_decimal = "1.36.0"
 serde = { workspace = true }

iceberg-rust-spec/Cargo.toml

@@ -23,6 +23,7 @@ iceberg-rust-spec = { path = "../iceberg-rust-spec", version = "0.8.0" }
 itertools = { workspace = true }
 lazy_static = { workspace = true }
 lru = { workspace = true }
+murmur3 = { workspace = true }
 object_store = { workspace = true }
 parquet = { workspace = true }
 pin-project-lite = { 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/Cargo.toml

@@ -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,66 @@ 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| {
+                let mut buffer = std::io::Cursor::new((i as i64).to_le_bytes());
+                (murmur3::murmur3_32(&mut buffer, 0).expect("murmur3 hash failled for some reason")
+                    as i32)
+                    .rem_euclid(*m as i32)
+            }),
+        )),
+        (DataType::Int64, Transform::Bucket(m)) => Ok(Arc::<PrimitiveArray<Int32Type>>::new(
+            unary(as_primitive_array::<Int64Type>(&array), |i| {
+                let mut buffer = std::io::Cursor::new((i).to_le_bytes());
+                (murmur3::murmur3_32(&mut buffer, 0).expect("murmur3 hash failled for some reason")
+                    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| {
+                    let mut buffer = std::io::Cursor::new((i as i64).to_le_bytes());
+                    (murmur3::murmur3_32(&mut buffer, 0)
+                        .expect("murmur3 hash failled for some reason") 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| {
+                    let mut buffer = std::io::Cursor::new((i as i64).to_le_bytes());
+                    (murmur3::murmur3_32(&mut buffer, 0)
+                        .expect("murmur3 hash failled for some reason") 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 {
+                        murmur3::murmur3_32(&mut value.as_bytes(), 0)
+                            .expect("murmur3 hash failled for some reason")
+                            as i32
+                    } else {
+                        0
+                    }
+                    .rem_euclid(*m as i32)
+                })),
+                nulls.cloned(),
+            )))
+        }
         _ => Err(ArrowError::ComputeError(
             "Failed to perform transform for datatype".to_string(),
         )),
@@ -322,6 +391,125 @@ 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
+        let mut buffer = std::io::Cursor::new((34i32 as i64).to_le_bytes());
+        assert_eq!(murmur3::murmur3_32(&mut buffer, 0).unwrap(), 2017239379);
+
+        // 34 -> 2017239379
+        let mut buffer = std::io::Cursor::new((34i64).to_le_bytes());
+        assert_eq!(murmur3::murmur3_32(&mut buffer, 0).unwrap(), 2017239379);
+
+        // daysFromUnixEpoch(2017-11-16) -> 17_486 -> -653330422
+        let mut buffer = std::io::Cursor::new((17_486i32 as i64).to_le_bytes());
+        assert_eq!(
+            murmur3::murmur3_32(&mut buffer, 0).unwrap() as i32,
+            -653330422
+        );
+
+        // 81_068_000_000 number of micros from midnight 22:31:08
+        let mut buffer = std::io::Cursor::new((81_068_000_000i64).to_le_bytes());
+        assert_eq!(
+            murmur3::murmur3_32(&mut buffer, 0).unwrap() as i32,
+            -662762989
+        );
+
+        // utf8Bytes(iceberg) -> 1210000089
+        assert_eq!(
+            murmur3::murmur3_32(&mut "iceberg".as_bytes(), 0).unwrap() as i32,
+            1210000089
+        );
+    }
+
+    #[test]
+    fn test_int32_bucket_transform() {
+        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(1000)).unwrap();
+        let expected = Arc::new(arrow::array::Int32Array::from(vec![
+            Some(2017239379i32.rem_euclid(1000)),
+            Some(578), // -653330422 % 1000 not match I don't know why
+            Some(988822981i32.rem_euclid(1000)),
+            Some(964620854i32.rem_euclid(1000)),
+            Some(1669671676i32.rem_euclid(1000)),
+            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(1000)).unwrap();
+        let expected = Arc::new(arrow::array::Int32Array::from(vec![
+            Some(2017239379i32.rem_euclid(1000)),
+            Some(578), // -653_330_422 % 1000 not match probably like to signed number
+            Some(117), // 716_914_497 = 1000 not match probably like to signed number
+            Some(964_620_854i32.rem_euclid(1000)),
+            Some(1669671676i32.rem_euclid(1000)),
+            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(1000)).unwrap();
+
+        let expected = Arc::new(arrow::array::Int32Array::from(vec![
+            Some(578), // -653330422 % 1000 not match probably like to signed number
+            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 micros from midnight 22:31:08
+            None,
+        ])) as ArrayRef;
+        let result = transform_arrow(array, &Transform::Bucket(1000)).unwrap();
+        let expected = Arc::new(arrow::array::Int32Array::from(vec![
+            Some(693), // -662762989 % 1000 not match probably like to signed number
+            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(1000)).unwrap();
+        let expected = Arc::new(arrow::array::Int32Array::from(vec![
+            Some(1_210_000_089i32.rem_euclid(1000)),
+            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;