Merge branch 'main' into issue-9645-ree-optimization

Author: liamzwbao

.github/workflows/take.yml

@@ -28,7 +28,7 @@ jobs:
     if: (!github.event.issue.pull_request) && github.event.comment.body == 'take'
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/github-script@v8
+      - uses: actions/github-script@v9
         with:
           script: |
             github.rest.issues.addAssignees({

arrow-array/Cargo.toml

@@ -49,7 +49,7 @@ num-complex = { version = "0.4.6", default-features = false, features = ["std"]
 num-integer = { version = "0.1.46", default-features = false, features = ["std"] }
 num-traits = { version = "0.2.19", default-features = false, features = ["std"] }
 half = { version = "2.1", default-features = false, features = ["num-traits"] }
-hashbrown = { version = "0.16.0", default-features = false }
+hashbrown = { version = "0.17.0", default-features = false }
 
 [package.metadata.docs.rs]
 all-features = true

arrow-array/src/builder/fixed_size_binary_dictionary_builder.rs

@@ -201,6 +201,10 @@ where
     fn finish_cloned(&self) -> ArrayRef {
         Arc::new(self.finish_cloned())
     }
+
+    fn finish_preserve_values(&mut self) -> ArrayRef {
+        Arc::new(self.finish_preserve_values())
+    }
 }
 
 impl<K> FixedSizeBinaryDictionaryBuilder<K>

arrow-array/src/builder/fixed_size_list_builder.rs

@@ -140,6 +140,10 @@ where
     fn finish_cloned(&self) -> ArrayRef {
         Arc::new(self.finish_cloned())
     }
+
+    fn finish_preserve_values(&mut self) -> ArrayRef {
+        Arc::new(self.finish_preserve_values())
+    }
 }
 
 impl<T: ArrayBuilder> FixedSizeListBuilder<T>
@@ -211,6 +215,28 @@ where
         FixedSizeListArray::new(field, self.list_len, values, nulls)
     }
 
+    fn finish_preserve_values(&mut self) -> FixedSizeListArray {
+        let len = self.len();
+        let values = self.values_builder.finish_preserve_values();
+        let nulls = self.null_buffer_builder.finish();
+
+        assert_eq!(
+            values.len(),
+            len * self.list_len as usize,
+            "Length of the child array ({}) must be the multiple of the value length ({}) and the array length ({}).",
+            values.len(),
+            self.list_len,
+            len,
+        );
+
+        let field = self
+            .field
+            .clone()
+            .unwrap_or_else(|| Arc::new(Field::new_list_field(values.data_type().clone(), true)));
+
+        FixedSizeListArray::new(field, self.list_len, values, nulls)
+    }
+
     /// Returns the current null buffer as a slice
     pub fn validity_slice(&self) -> Option<&[u8]> {
         self.null_buffer_builder.as_slice()
@@ -224,7 +250,7 @@ mod tests {
 
     use crate::Array;
     use crate::Int32Array;
-    use crate::builder::Int32Builder;
+    use crate::builder::{Int32Builder, tests::PreserveValuesMock};
 
     fn make_list_builder(
         include_null_element: bool,
@@ -491,4 +517,18 @@ mod tests {
 
         builder.finish();
     }
+
+    #[test]
+    fn test_finish_preserve_values() {
+        let mut builder = FixedSizeListBuilder::new(PreserveValuesMock::default(), 2);
+
+        builder.values().inner.append_value(0);
+        builder.values().inner.append_value(1);
+        builder.append(true);
+
+        let arr = builder.finish_preserve_values();
+
+        assert_eq!(1, arr.len());
+        assert_eq!(1, builder.values().called);
+    }
 }

arrow-array/src/builder/generic_bytes_dictionary_builder.rs

@@ -256,6 +256,10 @@ where
     fn finish_cloned(&self) -> ArrayRef {
         Arc::new(self.finish_cloned())
     }
+
+    fn finish_preserve_values(&mut self) -> ArrayRef {
+        Arc::new(self.finish_preserve_values())
+    }
 }
 
 impl<K, T> GenericByteDictionaryBuilder<K, T>

arrow-array/src/builder/generic_list_builder.rs

@@ -166,6 +166,10 @@ where
     fn finish_cloned(&self) -> ArrayRef {
         Arc::new(self.finish_cloned())
     }
+
+    fn finish_preserve_values(&mut self) -> ArrayRef {
+        Arc::new(self.finish_preserve_values())
+    }
 }
 
 impl<OffsetSize: OffsetSizeTrait, T: ArrayBuilder> GenericListBuilder<OffsetSize, T>
@@ -329,6 +333,23 @@ where
         GenericListArray::new(field, offsets, values, nulls)
     }
 
+    fn finish_preserve_values(&mut self) -> GenericListArray<OffsetSize> {
+        let values = self.values_builder.finish_preserve_values();
+        let nulls = self.null_buffer_builder.finish();
+
+        let offsets = Buffer::from_vec(std::mem::take(&mut self.offsets_builder));
+        // Safety: Safe by construction
+        let offsets = unsafe { OffsetBuffer::new_unchecked(offsets.into()) };
+        self.offsets_builder.push(OffsetSize::zero());
+
+        let field = match &self.field {
+            Some(f) => f.clone(),
+            None => Arc::new(Field::new_list_field(values.data_type().clone(), true)),
+        };
+
+        GenericListArray::new(field, offsets, values, nulls)
+    }
+
     /// Returns the current offsets buffer as a slice
     pub fn offsets_slice(&self) -> &[OffsetSize] {
         self.offsets_builder.as_slice()
@@ -364,7 +385,7 @@ where
 mod tests {
     use super::*;
     use crate::Int32Array;
-    use crate::builder::{Int32Builder, ListBuilder, make_builder};
+    use crate::builder::{Int32Builder, ListBuilder, make_builder, tests::PreserveValuesMock};
     use crate::cast::AsArray;
     use crate::types::Int32Type;
     use arrow_schema::DataType;
@@ -818,4 +839,17 @@ mod tests {
         builder.append_value([Some(1)]);
         builder.finish();
     }
+
+    #[test]
+    fn test_finish_preserve_values() {
+        let mut builder = ListBuilder::new(PreserveValuesMock::default());
+
+        builder.values().inner.append_value(1);
+        builder.append(true);
+
+        let arr = builder.finish_preserve_values();
+
+        assert_eq!(1, arr.len());
+        assert_eq!(1, builder.values().called);
+    }
 }

arrow-array/src/builder/generic_list_view_builder.rs

@@ -71,6 +71,10 @@ impl<OffsetSize: OffsetSizeTrait, T: ArrayBuilder> ArrayBuilder
     fn finish_cloned(&self) -> ArrayRef {
         Arc::new(self.finish_cloned())
     }
+
+    fn finish_preserve_values(&mut self) -> ArrayRef {
+        Arc::new(self.finish_preserve_values())
+    }
 }
 
 impl<OffsetSize: OffsetSizeTrait, T: ArrayBuilder> GenericListViewBuilder<OffsetSize, T> {
@@ -216,6 +220,23 @@ where
         GenericListViewArray::new(field, offsets, sizes, values, nulls)
     }
 
+    fn finish_preserve_values(&mut self) -> GenericListViewArray<OffsetSize> {
+        let values = self.values_builder.finish_preserve_values();
+        let nulls = self.null_buffer_builder.finish();
+        let offsets = Buffer::from_vec(std::mem::take(&mut self.offsets_builder));
+        self.current_offset = OffsetSize::zero();
+
+        // Safety: Safe by construction
+        let offsets = ScalarBuffer::from(offsets);
+        let sizes = Buffer::from_vec(std::mem::take(&mut self.sizes_builder));
+        let sizes = ScalarBuffer::from(sizes);
+        let field = match &self.field {
+            Some(f) => f.clone(),
+            None => Arc::new(Field::new("item", values.data_type().clone(), true)),
+        };
+        GenericListViewArray::new(field, offsets, sizes, values, nulls)
+    }
+
     /// Returns the current offsets buffer as a slice
     pub fn offsets_slice(&self) -> &[OffsetSize] {
         self.offsets_builder.as_slice()
@@ -245,7 +266,7 @@ where
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::builder::{Int32Builder, ListViewBuilder, make_builder};
+    use crate::builder::{Int32Builder, ListViewBuilder, make_builder, tests::PreserveValuesMock};
     use crate::cast::AsArray;
     use crate::types::Int32Type;
     use crate::{Array, Int32Array};
@@ -703,4 +724,17 @@ mod tests {
         builder.append_value([Some(1)]);
         builder.finish();
     }
+
+    #[test]
+    fn test_finish_preserve_values() {
+        let mut builder = ListViewBuilder::new(PreserveValuesMock::default());
+
+        builder.values().inner.append_value(1);
+        builder.append(true);
+
+        let arr = builder.finish_preserve_values();
+
+        assert_eq!(1, arr.len());
+        assert_eq!(1, builder.values().called);
+    }
 }

arrow-array/src/builder/map_builder.rs

@@ -214,6 +214,18 @@ impl<K: ArrayBuilder, V: ArrayBuilder> MapBuilder<K, V> {
         self.finish_helper(keys_arr, values_arr, offset_buffer, nulls, len)
     }
 
+    fn finish_preserve_values(&mut self) -> MapArray {
+        let len = self.len();
+        // Build the keys
+        let keys_arr = self.key_builder.finish_preserve_values();
+        let values_arr = self.value_builder.finish_preserve_values();
+        let offset_buffer = Buffer::from_vec(std::mem::take(&mut self.offsets_builder));
+        self.offsets_builder.push(0);
+        let null_bit_buffer = self.null_buffer_builder.finish();
+
+        self.finish_helper(keys_arr, values_arr, offset_buffer, null_bit_buffer, len)
+    }
+
     fn finish_helper(
         &self,
         keys_arr: Arc<dyn Array>,
@@ -287,6 +299,10 @@ impl<K: ArrayBuilder, V: ArrayBuilder> ArrayBuilder for MapBuilder<K, V> {
         Arc::new(self.finish_cloned())
     }
 
+    fn finish_preserve_values(&mut self) -> ArrayRef {
+        Arc::new(self.finish_preserve_values())
+    }
+
     fn as_any(&self) -> &dyn Any {
         self
     }
@@ -303,7 +319,7 @@ impl<K: ArrayBuilder, V: ArrayBuilder> ArrayBuilder for MapBuilder<K, V> {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::builder::{Int32Builder, StringBuilder, make_builder};
+    use crate::builder::{Int32Builder, StringBuilder, make_builder, tests::PreserveValuesMock};
     use crate::{Int32Array, StringArray};
     use std::collections::HashMap;
 
@@ -516,4 +532,23 @@ mod tests {
 
         builder.finish();
     }
+
+    #[test]
+    fn test_finish_preserve_values() {
+        let mut builder = MapBuilder::new(
+            None,
+            PreserveValuesMock::default(),
+            PreserveValuesMock::default(),
+        );
+
+        builder.keys().inner.append_value(1);
+        builder.values().inner.append_value(2);
+        builder.append(true).unwrap();
+
+        let map = builder.finish_preserve_values();
+
+        assert_eq!(1, map.len());
+        assert_eq!(1, builder.keys().called);
+        assert_eq!(1, builder.values().called);
+    }
 }

arrow-array/src/builder/mod.rs

@@ -341,6 +341,18 @@ pub trait ArrayBuilder: Any + Send + Sync {
     /// Builds the array without resetting the underlying builder.
     fn finish_cloned(&self) -> ArrayRef;
 
+    /// Builds the array without resetting the values builder.
+    ///
+    /// This is relevant for dictionary builders but also for composite builders.
+    /// Those are not affected directly, but will call the corresponding method
+    /// on their constituent builders.
+    ///
+    /// The default implementation just calls [`finish`][Self::finish] which is sufficient
+    /// for all but the above mentioned builders.
+    fn finish_preserve_values(&mut self) -> ArrayRef {
+        self.finish()
+    }
+
     /// Returns the builder as a non-mutable `Any` reference.
     ///
     /// This is most useful when one wants to call non-mutable APIs on a specific builder
@@ -376,6 +388,10 @@ impl ArrayBuilder for Box<dyn ArrayBuilder> {
         (**self).finish_cloned()
     }
 
+    fn finish_preserve_values(&mut self) -> ArrayRef {
+        (**self).finish_preserve_values()
+    }
+
     fn as_any(&self) -> &dyn Any {
         (**self).as_any()
     }
@@ -613,3 +629,45 @@ pub fn make_builder(datatype: &DataType, capacity: usize) -> Box<dyn ArrayBuilde
         t => unimplemented!("Data type {t} is not currently supported"),
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[derive(Default)]
+    pub struct PreserveValuesMock {
+        pub called: usize,
+        pub inner: Int32Builder,
+    }
+
+    impl ArrayBuilder for PreserveValuesMock {
+        fn as_any(&self) -> &dyn Any {
+            self
+        }
+
+        fn as_any_mut(&mut self) -> &mut dyn Any {
+            self
+        }
+
+        fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
+            self
+        }
+
+        fn len(&self) -> usize {
+            self.inner.len()
+        }
+
+        fn finish(&mut self) -> ArrayRef {
+            panic!("finish should never be called on PreserveValuesMock")
+        }
+
+        fn finish_cloned(&self) -> ArrayRef {
+            panic!("finish_cloned should never be called on PreserveValuesMock")
+        }
+
+        fn finish_preserve_values(&mut self) -> ArrayRef {
+            self.called += 1;
+            self.inner.finish_preserve_values()
+        }
+    }
+}

arrow-array/src/builder/primitive_dictionary_builder.rs

@@ -268,6 +268,10 @@ where
     fn finish_cloned(&self) -> ArrayRef {
         Arc::new(self.finish_cloned())
     }
+
+    fn finish_preserve_values(&mut self) -> ArrayRef {
+        Arc::new(self.finish_preserve_values())
+    }
 }
 
 impl<K, V> PrimitiveDictionaryBuilder<K, V>