Merge branch 'main' into feat/rg-reorder-by-statistics

Author: zhuqi-lucas

datafusion/functions/src/string/repeat.rs

@@ -15,13 +15,12 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::sync::Arc;
-
-use crate::utils::utf8_to_str_type;
-use arrow::array::{
-    Array, ArrayRef, AsArray, GenericStringArray, GenericStringBuilder, Int64Array,
-    StringArrayType, StringLikeArrayBuilder, StringViewArray, StringViewBuilder,
+use crate::strings::{
+    BulkNullStringArrayBuilder, GenericStringArrayBuilder, StringViewArrayBuilder,
 };
+use crate::utils::utf8_to_str_type;
+use arrow::array::{Array, ArrayRef, AsArray, Int64Array, StringArrayType};
+use arrow::buffer::NullBuffer;
 use arrow::datatypes::DataType;
 use arrow::datatypes::DataType::{LargeUtf8, Utf8, Utf8View};
 use datafusion_common::cast::as_int64_array;
@@ -190,43 +189,28 @@ fn repeat(string_array: &ArrayRef, count_array: &ArrayRef) -> Result<ArrayRef> {
                 number_array,
                 i32::MAX as usize,
             )?;
-            let builder = StringViewBuilder::with_capacity(string_array.len());
-            repeat_impl::<&StringViewArray, StringViewBuilder>(
-                &string_view_array,
-                number_array,
-                max_item_capacity,
-                builder,
-            )
+            let builder = StringViewArrayBuilder::with_capacity(string_array.len());
+            repeat_impl(&string_view_array, number_array, max_item_capacity, builder)
         }
         Utf8 => {
             let string_arr = string_array.as_string::<i32>();
             let (total_capacity, max_item_capacity) =
                 calculate_capacities(&string_arr, number_array, i32::MAX as usize)?;
-            let builder = GenericStringBuilder::<i32>::with_capacity(
+            let builder = GenericStringArrayBuilder::<i32>::with_capacity(
                 string_array.len(),
                 total_capacity,
             );
-            repeat_impl::<&GenericStringArray<i32>, GenericStringBuilder<i32>>(
-                &string_arr,
-                number_array,
-                max_item_capacity,
-                builder,
-            )
+            repeat_impl(&string_arr, number_array, max_item_capacity, builder)
         }
         LargeUtf8 => {
             let string_arr = string_array.as_string::<i64>();
             let (total_capacity, max_item_capacity) =
                 calculate_capacities(&string_arr, number_array, i64::MAX as usize)?;
-            let builder = GenericStringBuilder::<i64>::with_capacity(
+            let builder = GenericStringArrayBuilder::<i64>::with_capacity(
                 string_array.len(),
                 total_capacity,
             );
-            repeat_impl::<&GenericStringArray<i64>, GenericStringBuilder<i64>>(
-                &string_arr,
-                number_array,
-                max_item_capacity,
-                builder,
-            )
+            repeat_impl(&string_arr, number_array, max_item_capacity, builder)
         }
         other => exec_err!(
             "Unsupported data type {other:?} for function repeat. \
@@ -278,7 +262,7 @@ fn repeat_impl<'a, S, B>(
 ) -> Result<ArrayRef>
 where
     S: StringArrayType<'a> + 'a,
-    B: StringLikeArrayBuilder,
+    B: BulkNullStringArrayBuilder,
 {
     // Reusable buffer to avoid allocations in string.repeat()
     let mut buffer = Vec::<u8>::with_capacity(max_item_capacity);
@@ -301,12 +285,18 @@ where
         }
     }
 
-    // Fast path: no nulls in either array
-    if string_array.null_count() == 0 && number_array.null_count() == 0 {
+    // Output is null IFF either input is null
+    let nulls = NullBuffer::union(string_array.nulls(), number_array.nulls());
+
+    if let Some(ref n) = nulls {
         for i in 0..string_array.len() {
-            // SAFETY: i is within bounds (0..len) and null_count() == 0 guarantees valid value
+            if n.is_null(i) {
+                builder.append_placeholder();
+                continue;
+            }
+            // SAFETY: index `i` in both arrays is valid
             let string = unsafe { string_array.value_unchecked(i) };
-            let count = number_array.value(i);
+            let count = unsafe { number_array.value_unchecked(i) };
             if count > 0 {
                 repeat_to_buffer(&mut buffer, string, count as usize);
                 // SAFETY: buffer contains valid UTF-8 since we only copy from a valid &str
@@ -316,27 +306,30 @@ where
             }
         }
     } else {
-        // Slow path: handle nulls
-        for (string, number) in string_array.iter().zip(number_array.iter()) {
-            match (string, number) {
-                (Some(string), Some(count)) if count > 0 => {
-                    repeat_to_buffer(&mut buffer, string, count as usize);
-                    // SAFETY: buffer contains valid UTF-8 since we only copy from a valid &str
-                    builder
-                        .append_value(unsafe { std::str::from_utf8_unchecked(&buffer) });
-                }
-                (Some(_), Some(_)) => builder.append_value(""),
-                _ => builder.append_null(),
+        for i in 0..string_array.len() {
+            // SAFETY: no nulls, so every index in both arrays is valid
+            let string = unsafe { string_array.value_unchecked(i) };
+            let count = unsafe { number_array.value_unchecked(i) };
+            if count > 0 {
+                repeat_to_buffer(&mut buffer, string, count as usize);
+                // SAFETY: buffer contains valid UTF-8 since we only copy from a valid &str
+                builder.append_value(unsafe { std::str::from_utf8_unchecked(&buffer) });
+            } else {
+                builder.append_value("");
             }
         }
     }
 
-    Ok(Arc::new(builder.finish()) as ArrayRef)
+    builder.finish(nulls)
 }
 
 #[cfg(test)]
 mod tests {
-    use arrow::array::{Array, LargeStringArray, StringArray, StringViewArray};
+    use std::sync::Arc;
+
+    use arrow::array::{
+        Array, ArrayRef, Int64Array, LargeStringArray, StringArray, StringViewArray,
+    };
     use arrow::datatypes::DataType::{LargeUtf8, Utf8, Utf8View};
 
     use datafusion_common::ScalarValue;
@@ -444,4 +437,69 @@ mod tests {
 
         Ok(())
     }
+
+    // Slicing the input arrays produces a NullBuffer with a non-zero offset.
+    // The tests below use 6-row inputs sliced to (1, 4) so that:
+    //   slot 0 (orig 1): "a"  × 3    → "aaa"
+    //   slot 1 (orig 2): "bb" × 2    → "bbbb"
+    //   slot 2 (orig 3): "c"  × NULL → NULL (count-side null)
+    //   slot 3 (orig 4): NULL × 1    → NULL (string-side null)
+    fn sliced_offset_inputs<F>(make_strings: F) -> (ArrayRef, ArrayRef)
+    where
+        F: FnOnce(Vec<Option<&'static str>>) -> ArrayRef,
+    {
+        let strings = make_strings(vec![
+            None,
+            Some("a"),
+            Some("bb"),
+            Some("c"),
+            None,
+            Some("d"),
+        ]);
+        let counts: ArrayRef = Arc::new(Int64Array::from(vec![
+            Some(2),
+            Some(3),
+            Some(2),
+            None,
+            Some(1),
+            Some(2),
+        ]));
+        (strings.slice(1, 4), counts.slice(1, 4))
+    }
+
+    fn assert_sliced_offset_output<A: Array + 'static>(result: ArrayRef)
+    where
+        for<'a> &'a A: arrow::array::ArrayAccessor<Item = &'a str>,
+    {
+        let result = result.as_any().downcast_ref::<A>().unwrap();
+        assert_eq!(result.len(), 4);
+        assert_eq!(arrow::array::ArrayAccessor::value(&result, 0), "aaa");
+        assert_eq!(arrow::array::ArrayAccessor::value(&result, 1), "bbbb");
+        assert!(result.is_null(2));
+        assert!(result.is_null(3));
+        assert_eq!(result.null_count(), 2);
+    }
+
+    #[test]
+    fn test_repeat_sliced_string_with_null_offset() {
+        let (strings, counts) = sliced_offset_inputs(|v| Arc::new(StringArray::from(v)));
+        let result = super::repeat(&strings, &counts).unwrap();
+        assert_sliced_offset_output::<StringArray>(result);
+    }
+
+    #[test]
+    fn test_repeat_sliced_large_string_with_null_offset() {
+        let (strings, counts) =
+            sliced_offset_inputs(|v| Arc::new(LargeStringArray::from(v)));
+        let result = super::repeat(&strings, &counts).unwrap();
+        assert_sliced_offset_output::<LargeStringArray>(result);
+    }
+
+    #[test]
+    fn test_repeat_sliced_string_view_with_null_offset() {
+        let (strings, counts) =
+            sliced_offset_inputs(|v| Arc::new(StringViewArray::from(v)));
+        let result = super::repeat(&strings, &counts).unwrap();
+        assert_sliced_offset_output::<StringViewArray>(result);
+    }
 }

datafusion/functions/src/strings.rs

@@ -17,12 +17,14 @@
 
 use std::marker::PhantomData;
 use std::mem::size_of;
+use std::sync::Arc;
 
 use datafusion_common::{Result, exec_datafusion_err, internal_err};
 
 use arrow::array::{
-    Array, ArrayAccessor, ArrayDataBuilder, BinaryArray, ByteView, GenericStringArray,
-    LargeStringArray, OffsetSizeTrait, StringArray, StringViewArray, make_view,
+    Array, ArrayAccessor, ArrayDataBuilder, ArrayRef, BinaryArray, ByteView,
+    GenericStringArray, LargeStringArray, OffsetSizeTrait, StringArray, StringViewArray,
+    make_view,
 };
 use arrow::buffer::{Buffer, MutableBuffer, NullBuffer, ScalarBuffer};
 use arrow::datatypes::DataType;
@@ -473,6 +475,7 @@ impl<O: OffsetSizeTrait> GenericStringArrayBuilder<O> {
     /// # Panics
     ///
     /// Panics if the cumulative byte length exceeds `O::MAX`.
+    #[inline]
     pub fn append_value(&mut self, value: &str) {
         self.value_buffer.extend_from_slice(value.as_bytes());
         let next_offset =
@@ -482,6 +485,7 @@ impl<O: OffsetSizeTrait> GenericStringArrayBuilder<O> {
 
     /// Append an empty placeholder row. The corresponding slot must be masked
     /// as null by the null buffer passed to `finish`.
+    #[inline]
     pub fn append_placeholder(&mut self) {
         let next_offset =
             O::from_usize(self.value_buffer.len()).expect("byte array offset overflow");
@@ -672,6 +676,47 @@ impl StringViewArrayBuilder {
     }
 }
 
+/// Trait abstracting over the bulk-NULL string array builders.
+///
+/// Similar to Arrow's `StringLikeArrayBuilder`, this allows generic dispatch
+/// over the three string array types (Utf8, LargeUtf8, Utf8View) when the
+/// function body is uniform across them.
+pub(crate) trait BulkNullStringArrayBuilder {
+    fn append_value(&mut self, value: &str);
+    fn append_placeholder(&mut self);
+    fn finish(self, nulls: Option<NullBuffer>) -> Result<ArrayRef>;
+}
+
+impl<O: OffsetSizeTrait> BulkNullStringArrayBuilder for GenericStringArrayBuilder<O> {
+    #[inline]
+    fn append_value(&mut self, value: &str) {
+        GenericStringArrayBuilder::<O>::append_value(self, value)
+    }
+    #[inline]
+    fn append_placeholder(&mut self) {
+        GenericStringArrayBuilder::<O>::append_placeholder(self)
+    }
+    fn finish(self, nulls: Option<NullBuffer>) -> Result<ArrayRef> {
+        Ok(Arc::new(GenericStringArrayBuilder::<O>::finish(
+            self, nulls,
+        )?))
+    }
+}
+
+impl BulkNullStringArrayBuilder for StringViewArrayBuilder {
+    #[inline]
+    fn append_value(&mut self, value: &str) {
+        StringViewArrayBuilder::append_value(self, value)
+    }
+    #[inline]
+    fn append_placeholder(&mut self) {
+        StringViewArrayBuilder::append_placeholder(self)
+    }
+    fn finish(self, nulls: Option<NullBuffer>) -> Result<ArrayRef> {
+        Ok(Arc::new(StringViewArrayBuilder::finish(self, nulls)?))
+    }
+}
+
 /// Append a new view to the views buffer with the given substr.
 ///
 /// Callers are responsible for their own null tracking.
@@ -962,4 +1007,65 @@ mod tests {
             assert_eq!(array.value(i), value);
         }
     }
+
+    /// Build an array via `BulkNullStringArrayBuilder` to verify that the
+    /// trait methods produce the same result as the inherent methods.
+    fn build_via_trait<B: BulkNullStringArrayBuilder>(
+        mut builder: B,
+        nulls: Option<NullBuffer>,
+    ) -> ArrayRef {
+        builder.append_value("a");
+        builder.append_placeholder();
+        builder.append_value("hello world!");
+        builder.finish(nulls).unwrap()
+    }
+
+    #[test]
+    fn bulk_null_trait_string_i32() {
+        let builder = GenericStringArrayBuilder::<i32>::with_capacity(3, 16);
+        let nulls = NullBuffer::from(vec![true, false, true]);
+        let array = build_via_trait(builder, Some(nulls));
+        let array = array.as_any().downcast_ref::<StringArray>().unwrap();
+        assert_eq!(array.len(), 3);
+        assert_eq!(array.value(0), "a");
+        assert!(array.is_null(1));
+        assert_eq!(array.value(2), "hello world!");
+    }
+
+    #[test]
+    fn bulk_null_trait_string_i64() {
+        let builder = GenericStringArrayBuilder::<i64>::with_capacity(3, 16);
+        let nulls = NullBuffer::from(vec![true, false, true]);
+        let array = build_via_trait(builder, Some(nulls));
+        let array = array.as_any().downcast_ref::<LargeStringArray>().unwrap();
+        assert_eq!(array.len(), 3);
+        assert_eq!(array.value(0), "a");
+        assert!(array.is_null(1));
+        assert_eq!(array.value(2), "hello world!");
+    }
+
+    #[test]
+    fn bulk_null_trait_string_view() {
+        let builder = StringViewArrayBuilder::with_capacity(3);
+        let nulls = NullBuffer::from(vec![true, false, true]);
+        let array = build_via_trait(builder, Some(nulls));
+        let array = array.as_any().downcast_ref::<StringViewArray>().unwrap();
+        assert_eq!(array.len(), 3);
+        assert_eq!(array.value(0), "a");
+        assert!(array.is_null(1));
+        assert_eq!(array.value(2), "hello world!");
+    }
+
+    #[test]
+    fn bulk_null_trait_no_nulls() {
+        let mut builder = GenericStringArrayBuilder::<i32>::with_capacity(2, 8);
+        BulkNullStringArrayBuilder::append_value(&mut builder, "x");
+        BulkNullStringArrayBuilder::append_value(&mut builder, "yy");
+        let array = BulkNullStringArrayBuilder::finish(builder, None).unwrap();
+        let array = array.as_any().downcast_ref::<StringArray>().unwrap();
+        assert_eq!(array.len(), 2);
+        assert_eq!(array.value(0), "x");
+        assert_eq!(array.value(1), "yy");
+        assert_eq!(array.null_count(), 0);
+    }
 }