## Coverage Report

#### Created: 2025-11-22 06:00

next uncovered line (L), next uncovered region (R), next uncovered branch (B)

    /Users/andrewlamb/Software/datafusion/datafusion/functions/src/strings.rs

    Line
    Count
    Source
    1
    // Licensed to the Apache Software Foundation (ASF) under one
    2
    // or more contributor license agreements.  See the NOTICE file
    3
    // distributed with this work for additional information
    4
    // regarding copyright ownership.  The ASF licenses this file
    5
    // to you under the Apache License, Version 2.0 (the
    6
    // "License"); you may not use this file except in compliance
    7
    // with the License.  You may obtain a copy of the License at
    8
    //
    9
    //   http://www.apache.org/licenses/LICENSE-2.0
    10
    //
    11
    // Unless required by applicable law or agreed to in writing,
    12
    // software distributed under the License is distributed on an
    13
    // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
    14
    // KIND, either express or implied.  See the License for the
    15
    // specific language governing permissions and limitations
    16
    // under the License.
    17
    18
    use std::mem::size_of;
    19
    20
    use arrow::array::{
    21
        make_view, Array, ArrayAccessor, ArrayDataBuilder, ByteView, LargeStringArray,
    22
        NullBufferBuilder, StringArray, StringViewArray, StringViewBuilder,
    23
    };
    24
    use arrow::buffer::{MutableBuffer, NullBuffer};
    25
    use arrow::datatypes::DataType;
    26
    27
    /// Optimized version of the StringBuilder in Arrow that:
    28
    /// 1. Precalculating the expected length of the result, avoiding reallocations.
    29
    /// 2. Avoids creating / incrementally creating a `NullBufferBuilder`
    30
    pub struct StringArrayBuilder {
    31
        offsets_buffer: MutableBuffer,
    32
        value_buffer: MutableBuffer,
    33
    }
    34
    35
    impl StringArrayBuilder {
    36
    0
        pub fn with_capacity(item_capacity: usize, data_capacity: usize) -> Self {
    37
    0
            let capacity = item_capacity
    38
    0
                .checked_add(1)
    39
    0
                .map(|i| i.saturating_mul(size_of::<i32>()))
    40
    0
                .expect("capacity integer overflow");
    41
    42
    0
            let mut offsets_buffer = MutableBuffer::with_capacity(capacity);
    43
            // SAFETY: the first offset value is definitely not going to exceed the bounds.
    44
    0
            unsafe { offsets_buffer.push_unchecked(0_i32) };
    45
    0
            Self {
    46
    0
                offsets_buffer,
    47
    0
                value_buffer: MutableBuffer::with_capacity(data_capacity),
    48
    0
            }
    49
    0
        }
    50
    51
    0
        pub fn write<const CHECK_VALID: bool>(
    52
    0
            &mut self,
    53
    0
            column: &ColumnarValueRef,
    54
    0
            i: usize,
    55
    0
        ) {
    56
    0
            match column {
    57
    0
                ColumnarValueRef::Scalar(s) => {
    58
    0
                    self.value_buffer.extend_from_slice(s);
    59
    0
                }
    60
    0
                ColumnarValueRef::NullableArray(array) => {
    61
    0
                    if !CHECK_VALID || array.is_valid(i) {
    62
    0
                        self.value_buffer
    63
    0
                            .extend_from_slice(array.value(i).as_bytes());
    64
    0
                    }
    65
                }
    66
    0
                ColumnarValueRef::NullableLargeStringArray(array) => {
    67
    0
                    if !CHECK_VALID || array.is_valid(i) {
    68
    0
                        self.value_buffer
    69
    0
                            .extend_from_slice(array.value(i).as_bytes());
    70
    0
                    }
    71
                }
    72
    0
                ColumnarValueRef::NullableStringViewArray(array) => {
    73
    0
                    if !CHECK_VALID || array.is_valid(i) {
    74
    0
                        self.value_buffer
    75
    0
                            .extend_from_slice(array.value(i).as_bytes());
    76
    0
                    }
    77
                }
    78
    0
                ColumnarValueRef::NonNullableArray(array) => {
    79
    0
                    self.value_buffer
    80
    0
                        .extend_from_slice(array.value(i).as_bytes());
    81
    0
                }
    82
    0
                ColumnarValueRef::NonNullableLargeStringArray(array) => {
    83
    0
                    self.value_buffer
    84
    0
                        .extend_from_slice(array.value(i).as_bytes());
    85
    0
                }
    86
    0
                ColumnarValueRef::NonNullableStringViewArray(array) => {
    87
    0
                    self.value_buffer
    88
    0
                        .extend_from_slice(array.value(i).as_bytes());
    89
    0
                }
    90
            }
    91
    0
        }
    92
    93
    0
        pub fn append_offset(&mut self) {
    94
    0
            let next_offset: i32 = self
    95
    0
                .value_buffer
    96
    0
                .len()
    97
    0
                .try_into()
    98
    0
                .expect("byte array offset overflow");
    99
    0
            self.offsets_buffer.push(next_offset);
    100
    0
        }
    101
    102
        /// Finalize the builder into a concrete [`StringArray`].
    103
        ///
    104
        /// # Panics
    105
        ///
    106
        /// This method can panic when:
    107
        ///
    108
        /// - the provided `null_buffer` is not the same length as the `offsets_buffer`.
    109
    0
        pub fn finish(self, null_buffer: Option<NullBuffer>) -> StringArray {
    110
    0
            let row_count = self.offsets_buffer.len() / size_of::<i32>() - 1;
    111
    0
            if let Some(ref null_buffer) = null_buffer {
    112
    0
                assert_eq!(
    113
    0
                    null_buffer.len(),
    114
                    row_count,
    115
    0
                    "Null buffer and offsets buffer must be the same length"
    116
                );
    117
    0
            }
    118
    0
            let array_builder = ArrayDataBuilder::new(DataType::Utf8)
    119
    0
                .len(row_count)
    120
    0
                .add_buffer(self.offsets_buffer.into())
    121
    0
                .add_buffer(self.value_buffer.into())
    122
    0
                .nulls(null_buffer);
    123
            // SAFETY: all data that was appended was valid UTF8 and the values
    124
            // and offsets were created correctly
    125
    0
            let array_data = unsafe { array_builder.build_unchecked() };
    126
    0
            StringArray::from(array_data)
    127
    0
        }
    128
    }
    129
    130
    pub struct StringViewArrayBuilder {
    131
        builder: StringViewBuilder,
    132
        block: String,
    133
    }
    134
    135
    impl StringViewArrayBuilder {
    136
    0
        pub fn with_capacity(_item_capacity: usize, data_capacity: usize) -> Self {
    137
    0
            let builder = StringViewBuilder::with_capacity(data_capacity);
    138
    0
            Self {
    139
    0
                builder,
    140
    0
                block: String::new(),
    141
    0
            }
    142
    0
        }
    143
    144
    0
        pub fn write<const CHECK_VALID: bool>(
    145
    0
            &mut self,
    146
    0
            column: &ColumnarValueRef,
    147
    0
            i: usize,
    148
    0
        ) {
    149
    0
            match column {
    150
    0
                ColumnarValueRef::Scalar(s) => {
    151
    0
                    self.block.push_str(std::str::from_utf8(s).unwrap());
    152
    0
                }
    153
    0
                ColumnarValueRef::NullableArray(array) => {
    154
    0
                    if !CHECK_VALID || array.is_valid(i) {
    155
    0
                        self.block.push_str(
    156
    0
                            std::str::from_utf8(array.value(i).as_bytes()).unwrap(),
    157
    0
                        );
    158
    0
                    }
    159
                }
    160
    0
                ColumnarValueRef::NullableLargeStringArray(array) => {
    161
    0
                    if !CHECK_VALID || array.is_valid(i) {
    162
    0
                        self.block.push_str(
    163
    0
                            std::str::from_utf8(array.value(i).as_bytes()).unwrap(),
    164
    0
                        );
    165
    0
                    }
    166
                }
    167
    0
                ColumnarValueRef::NullableStringViewArray(array) => {
    168
    0
                    if !CHECK_VALID || array.is_valid(i) {
    169
    0
                        self.block.push_str(
    170
    0
                            std::str::from_utf8(array.value(i).as_bytes()).unwrap(),
    171
    0
                        );
    172
    0
                    }
    173
                }
    174
    0
                ColumnarValueRef::NonNullableArray(array) => {
    175
    0
                    self.block
    176
    0
                        .push_str(std::str::from_utf8(array.value(i).as_bytes()).unwrap());
    177
    0
                }
    178
    0
                ColumnarValueRef::NonNullableLargeStringArray(array) => {
    179
    0
                    self.block
    180
    0
                        .push_str(std::str::from_utf8(array.value(i).as_bytes()).unwrap());
    181
    0
                }
    182
    0
                ColumnarValueRef::NonNullableStringViewArray(array) => {
    183
    0
                    self.block
    184
    0
                        .push_str(std::str::from_utf8(array.value(i).as_bytes()).unwrap());
    185
    0
                }
    186
            }
    187
    0
        }
    188
    189
    0
        pub fn append_offset(&mut self) {
    190
    0
            self.builder.append_value(&self.block);
    191
    0
            self.block = String::new();
    192
    0
        }
    193
    194
    0
        pub fn finish(mut self) -> StringViewArray {
    195
    0
            self.builder.finish()
    196
    0
        }
    197
    }
    198
    199
    pub struct LargeStringArrayBuilder {
    200
        offsets_buffer: MutableBuffer,
    201
        value_buffer: MutableBuffer,
    202
    }
    203
    204
    impl LargeStringArrayBuilder {
    205
    0
        pub fn with_capacity(item_capacity: usize, data_capacity: usize) -> Self {
    206
    0
            let capacity = item_capacity
    207
    0
                .checked_add(1)
    208
    0
                .map(|i| i.saturating_mul(size_of::<i64>()))
    209
    0
                .expect("capacity integer overflow");
    210
    211
    0
            let mut offsets_buffer = MutableBuffer::with_capacity(capacity);
    212
            // SAFETY: the first offset value is definitely not going to exceed the bounds.
    213
    0
            unsafe { offsets_buffer.push_unchecked(0_i64) };
    214
    0
            Self {
    215
    0
                offsets_buffer,
    216
    0
                value_buffer: MutableBuffer::with_capacity(data_capacity),
    217
    0
            }
    218
    0
        }
    219
    220
    0
        pub fn write<const CHECK_VALID: bool>(
    221
    0
            &mut self,
    222
    0
            column: &ColumnarValueRef,
    223
    0
            i: usize,
    224
    0
        ) {
    225
    0
            match column {
    226
    0
                ColumnarValueRef::Scalar(s) => {
    227
    0
                    self.value_buffer.extend_from_slice(s);
    228
    0
                }
    229
    0
                ColumnarValueRef::NullableArray(array) => {
    230
    0
                    if !CHECK_VALID || array.is_valid(i) {
    231
    0
                        self.value_buffer
    232
    0
                            .extend_from_slice(array.value(i).as_bytes());
    233
    0
                    }
    234
                }
    235
    0
                ColumnarValueRef::NullableLargeStringArray(array) => {
    236
    0
                    if !CHECK_VALID || array.is_valid(i) {
    237
    0
                        self.value_buffer
    238
    0
                            .extend_from_slice(array.value(i).as_bytes());
    239
    0
                    }
    240
                }
    241
    0
                ColumnarValueRef::NullableStringViewArray(array) => {
    242
    0
                    if !CHECK_VALID || array.is_valid(i) {
    243
    0
                        self.value_buffer
    244
    0
                            .extend_from_slice(array.value(i).as_bytes());
    245
    0
                    }
    246
                }
    247
    0
                ColumnarValueRef::NonNullableArray(array) => {
    248
    0
                    self.value_buffer
    249
    0
                        .extend_from_slice(array.value(i).as_bytes());
    250
    0
                }
    251
    0
                ColumnarValueRef::NonNullableLargeStringArray(array) => {
    252
    0
                    self.value_buffer
    253
    0
                        .extend_from_slice(array.value(i).as_bytes());
    254
    0
                }
    255
    0
                ColumnarValueRef::NonNullableStringViewArray(array) => {
    256
    0
                    self.value_buffer
    257
    0
                        .extend_from_slice(array.value(i).as_bytes());
    258
    0
                }
    259
            }
    260
    0
        }
    261
    262
    0
        pub fn append_offset(&mut self) {
    263
    0
            let next_offset: i64 = self
    264
    0
                .value_buffer
    265
    0
                .len()
    266
    0
                .try_into()
    267
    0
                .expect("byte array offset overflow");
    268
    0
            self.offsets_buffer.push(next_offset);
    269
    0
        }
    270
    271
        /// Finalize the builder into a concrete [`LargeStringArray`].
    272
        ///
    273
        /// # Panics
    274
        ///
    275
        /// This method can panic when:
    276
        ///
    277
        /// - the provided `null_buffer` is not the same length as the `offsets_buffer`.
    278
    0
        pub fn finish(self, null_buffer: Option<NullBuffer>) -> LargeStringArray {
    279
    0
            let row_count = self.offsets_buffer.len() / size_of::<i64>() - 1;
    280
    0
            if let Some(ref null_buffer) = null_buffer {
    281
    0
                assert_eq!(
    282
    0
                    null_buffer.len(),
    283
                    row_count,
    284
    0
                    "Null buffer and offsets buffer must be the same length"
    285
                );
    286
    0
            }
    287
    0
            let array_builder = ArrayDataBuilder::new(DataType::LargeUtf8)
    288
    0
                .len(row_count)
    289
    0
                .add_buffer(self.offsets_buffer.into())
    290
    0
                .add_buffer(self.value_buffer.into())
    291
    0
                .nulls(null_buffer);
    292
            // SAFETY: all data that was appended was valid Large UTF8 and the values
    293
            // and offsets were created correctly
    294
    0
            let array_data = unsafe { array_builder.build_unchecked() };
    295
    0
            LargeStringArray::from(array_data)
    296
    0
        }
    297
    }
    298
    299
    /// Append a new view to the views buffer with the given substr
    300
    ///
    301
    /// # Safety
    302
    ///
    303
    /// original_view must be a valid view (the format described on
    304
    /// [`GenericByteViewArray`](arrow::array::GenericByteViewArray).
    305
    ///
    306
    /// # Arguments
    307
    /// - views_buffer: The buffer to append the new view to
    308
    /// - null_builder: The buffer to append the null value to
    309
    /// - original_view: The original view value
    310
    /// - substr: The substring to append. Must be a valid substring of the original view
    311
    /// - start_offset: The start offset of the substring in the view
    312
    0
    pub fn make_and_append_view(
    313
    0
        views_buffer: &mut Vec<u128>,
    314
    0
        null_builder: &mut NullBufferBuilder,
    315
    0
        original_view: &u128,
    316
    0
        substr: &str,
    317
    0
        start_offset: u32,
    318
    0
    ) {
    319
    0
        let substr_len = substr.len();
    320
    0
        let sub_view = if substr_len > 12 {
    321
    0
            let view = ByteView::from(*original_view);
    322
    0
            make_view(
    323
    0
                substr.as_bytes(),
    324
    0
                view.buffer_index,
    325
    0
                view.offset + start_offset,
    326
            )
    327
        } else {
    328
            // inline value does not need block id or offset
    329
    0
            make_view(substr.as_bytes(), 0, 0)
    330
        };
    331
    0
        views_buffer.push(sub_view);
    332
    0
        null_builder.append_non_null();
    333
    0
    }
    334
    335
    #[derive(Debug)]
    336
    pub enum ColumnarValueRef<'a> {
    337
        Scalar(&'a [u8]),
    338
        NullableArray(&'a StringArray),
    339
        NonNullableArray(&'a StringArray),
    340
        NullableLargeStringArray(&'a LargeStringArray),
    341
        NonNullableLargeStringArray(&'a LargeStringArray),
    342
        NullableStringViewArray(&'a StringViewArray),
    343
        NonNullableStringViewArray(&'a StringViewArray),
    344
    }
    345
    346
    impl ColumnarValueRef<'_> {
    347
        #[inline]
    348
    0
        pub fn is_valid(&self, i: usize) -> bool {
    349
    0
            match &self {
    350
    0
                Self::Scalar(_)
    351
    0
                | Self::NonNullableArray(_)
    352
    0
                | Self::NonNullableLargeStringArray(_)
    353
    0
                | Self::NonNullableStringViewArray(_) => true,
    354
    0
                Self::NullableArray(array) => array.is_valid(i),
    355
    0
                Self::NullableStringViewArray(array) => array.is_valid(i),
    356
    0
                Self::NullableLargeStringArray(array) => array.is_valid(i),
    357
            }
    358
    0
        }
    359
    360
        #[inline]
    361
    0
        pub fn nulls(&self) -> Option<NullBuffer> {
    362
    0
            match &self {
    363
    0
                Self::Scalar(_)
    364
    0
                | Self::NonNullableArray(_)
    365
    0
                | Self::NonNullableStringViewArray(_)
    366
    0
                | Self::NonNullableLargeStringArray(_) => None,
    367
    0
                Self::NullableArray(array) => array.nulls().cloned(),
    368
    0
                Self::NullableStringViewArray(array) => array.nulls().cloned(),
    369
    0
                Self::NullableLargeStringArray(array) => array.nulls().cloned(),
    370
            }
    371
    0
        }
    372
    }
    373
    374
    #[cfg(test)]
    375
    mod tests {
    376
        use super::*;
    377
    378
        #[test]
    379
        #[should_panic(expected = "capacity integer overflow")]
    380
        fn test_overflow_string_array_builder() {
    381
            let _builder = StringArrayBuilder::with_capacity(usize::MAX, usize::MAX);
    382
        }
    383
    384
        #[test]
    385
        #[should_panic(expected = "capacity integer overflow")]
    386
        fn test_overflow_large_string_array_builder() {
    387
            let _builder = LargeStringArrayBuilder::with_capacity(usize::MAX, usize::MAX);
    388
        }
    389
    }