[Arrow]Configure max deduplication length when deduplicating strings while building the array

Author: lichuang

arrow-array/src/builder/generic_bytes_view_builder.rs

@@ -87,6 +87,7 @@ pub struct GenericByteViewBuilder<T: ByteViewType + ?Sized> {
     /// Some if deduplicating strings
     /// map `<string hash> -> <index to the views>`
     string_tracker: Option<(HashTable<usize>, ahash::RandomState)>,
+    max_deduplication_len: Option<u32>,
     phantom: PhantomData<T>,
 }
 
@@ -107,10 +108,25 @@ impl<T: ByteViewType + ?Sized> GenericByteViewBuilder<T> {
                 current_size: STARTING_BLOCK_SIZE,
             },
             string_tracker: None,
+            max_deduplication_len: None,
             phantom: Default::default(),
         }
     }
 
+    /// Configure max deduplication length when deduplicating strings while building the array.
+    /// Default is [`MAX_INLINE_VIEW_LEN`] bytes.
+    /// See <https://github.com/apache/arrow-rs/issues/7187> for more details on the implications.
+    pub fn with_max_deduplication_len(self, max_deduplication_len: u32) -> Self {
+        debug_assert!(
+            max_deduplication_len > 0,
+            "max_deduplication_len must be greater than 0"
+        );
+        Self {
+            max_deduplication_len: Some(max_deduplication_len),
+            ..self
+        }
+    }
+
     /// Set a fixed buffer size for variable length strings
     ///
     /// The block size is the size of the buffer used to store values greater
@@ -334,35 +350,41 @@ impl<T: ByteViewType + ?Sized> GenericByteViewBuilder<T> {
 
         // Deduplication if:
         // (1) deduplication is enabled.
-        // (2) len > 12
-        if let Some((mut ht, hasher)) = self.string_tracker.take() {
-            let hash_val = hasher.hash_one(v);
-            let hasher_fn = |v: &_| hasher.hash_one(v);
-
-            let entry = ht.entry(
-                hash_val,
-                |idx| {
-                    let stored_value = self.get_value(*idx);
-                    v == stored_value
-                },
-                hasher_fn,
-            );
-            match entry {
-                Entry::Occupied(occupied) => {
-                    // If the string already exists, we will directly use the view
-                    let idx = occupied.get();
-                    self.views_buffer.push(self.views_buffer[*idx]);
-                    self.null_buffer_builder.append_non_null();
-                    self.string_tracker = Some((ht, hasher));
-                    return Ok(());
-                }
-                Entry::Vacant(vacant) => {
-                    // o.w. we insert the (string hash -> view index)
-                    // the idx is current length of views_builder, as we are inserting a new view
-                    vacant.insert(self.views_buffer.len());
+        // (2) len > `MAX_INLINE_VIEW_LEN` and len < `max_deduplication_len`
+        let can_deduplicate = match self.max_deduplication_len {
+            Some(max_deduplication_len) => length <= max_deduplication_len,
+            None => true,
+        };
+        if can_deduplicate {
+            if let Some((mut ht, hasher)) = self.string_tracker.take() {
+                let hash_val = hasher.hash_one(v);
+                let hasher_fn = |v: &_| hasher.hash_one(v);
+
+                let entry = ht.entry(
+                    hash_val,
+                    |idx| {
+                        let stored_value = self.get_value(*idx);
+                        v == stored_value
+                    },
+                    hasher_fn,
+                );
+                match entry {
+                    Entry::Occupied(occupied) => {
+                        // If the string already exists, we will directly use the view
+                        let idx = occupied.get();
+                        self.views_buffer.push(self.views_buffer[*idx]);
+                        self.null_buffer_builder.append_non_null();
+                        self.string_tracker = Some((ht, hasher));
+                        return Ok(());
+                    }
+                    Entry::Vacant(vacant) => {
+                        // o.w. we insert the (string hash -> view index)
+                        // the idx is current length of views_builder, as we are inserting a new view
+                        vacant.insert(self.views_buffer.len());
+                    }
                 }
+                self.string_tracker = Some((ht, hasher));
             }
-            self.string_tracker = Some((ht, hasher));
         }
 
         let required_cap = self.in_progress.len() + v.len();
@@ -636,8 +658,59 @@ pub fn make_view(data: &[u8], block_id: u32, offset: u32) -> u128 {
 mod tests {
     use core::str;
 
+    use arrow_buffer::ArrowNativeType;
+
     use super::*;
 
+    #[test]
+    fn test_string_max_deduplication_len() {
+        let value_1 = "short";
+        let value_2 = "not so similar string but long";
+        let value_3 = "1234567890123";
+
+        let mut builder = StringViewBuilder::new()
+            .with_deduplicate_strings()
+            .with_max_deduplication_len(MAX_INLINE_VIEW_LEN * 2);
+        assert!(builder.string_tracker.is_some());
+        assert!(builder.max_deduplication_len.is_some());
+        // Since assert `max_deduplication_len` is some, so it is safe to unwrap here
+        let max_deduplication_len = builder.max_deduplication_len.unwrap();
+        assert!(max_deduplication_len > MAX_INLINE_VIEW_LEN);
+        assert!(value_1.len() < MAX_INLINE_VIEW_LEN.as_usize());
+        assert!(value_2.len() > max_deduplication_len.as_usize());
+        assert!(
+            value_3.len() > MAX_INLINE_VIEW_LEN.as_usize()
+                && value_3.len() < max_deduplication_len.as_usize()
+        );
+
+        let value_checker = |v: &[u8], builder: &StringViewBuilder| {
+            // safe to unwrap
+            let (ht, hasher) = builder.string_tracker.clone().unwrap();
+            let hash_val = hasher.hash_one(v);
+
+            ht.find(hash_val, |idx| {
+                let stored_value = builder.get_value(*idx);
+                v == stored_value
+            })
+            .cloned()
+        };
+
+        // append value1, it MUST not in the string_tracker
+        let v: &[u8] = value_1.as_ref();
+        builder.append_value(value_1);
+        assert!(value_checker(v, &builder).is_none());
+
+        // append value2, it MUST not in the string_tracker
+        let v: &[u8] = value_2.as_ref();
+        builder.append_value(value_2);
+        assert!(value_checker(v, &builder).is_none());
+
+        // append value3, it MUST in the string_tracker
+        let v: &[u8] = value_3.as_ref();
+        builder.append_value(value_3);
+        assert!(value_checker(v, &builder).is_some());
+    }
+
     #[test]
     fn test_string_view_deduplicate() {
         let value_1 = "long string to test string view";