Add Truncate for Binary type (#920)

apache · Jan 30, 2025 · e48b431 · e48b431
1 parent c5cc9e4
commit e48b431
Showing 1 changed file with 28 additions and 0 deletions.
diff --git a/crates/iceberg/src/transform/truncate.rs b/crates/iceberg/src/transform/truncate.rs
@@ -42,6 +42,11 @@ impl Truncate {
         }
     }
 
+    #[inline]
+    fn truncate_binary(s: &[u8], width: usize) -> &[u8] {
+        &s[0..width]
+    }
+
     #[inline]
     fn truncate_i32(v: i32, width: i32) -> i32 {
         v - v.rem_euclid(width)
@@ -119,6 +124,18 @@ impl TransformFunction for Truncate {
                 );
                 Ok(Arc::new(res))
             }
+            DataType::Binary => {
+                let len = self.width as usize;
+                let res: arrow_array::BinaryArray = arrow_array::BinaryArray::from_iter(
+                    input
+                        .as_any()
+                        .downcast_ref::<arrow_array::BinaryArray>()
+                        .unwrap()
+                        .iter()
+                        .map(|v| v.map(|v| Self::truncate_binary(v, len))),
+                );
+                Ok(Arc::new(res))
+            }
             _ => Err(crate::Error::new(
                 crate::ErrorKind::FeatureUnsupported,
                 format!(
@@ -747,6 +764,17 @@ mod test {
                 .value(0),
             "ice"
         );
+
+        // test binary
+        let input = Arc::new(arrow_array::BinaryArray::from_vec(vec![b"iceberg"]));
+        let res = super::Truncate::new(3).transform(input).unwrap();
+        assert_eq!(
+            res.as_any()
+                .downcast_ref::<arrow_array::BinaryArray>()
+                .unwrap()
+                .value(0),
+            b"ice"
+        );
     }
 
     #[test]