## 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/physical-expr-common/src/binary_view_map.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
    //! [`ArrowBytesViewMap`] and [`ArrowBytesViewSet`] for storing maps/sets of values from
    19
    //! `StringViewArray`/`BinaryViewArray`.
    20
    //! Much of the code is from `binary_map.rs`, but with simpler implementation because we directly use the
    21
    //! [`GenericByteViewBuilder`].
    22
    use crate::binary_map::OutputType;
    23
    use ahash::RandomState;
    24
    use arrow::array::cast::AsArray;
    25
    use arrow::array::{Array, ArrayBuilder, ArrayRef, GenericByteViewBuilder};
    26
    use arrow::datatypes::{BinaryViewType, ByteViewType, DataType, StringViewType};
    27
    use datafusion_common::hash_utils::create_hashes;
    28
    use datafusion_common::utils::proxy::{HashTableAllocExt, VecAllocExt};
    29
    use std::fmt::Debug;
    30
    use std::sync::Arc;
    31
    32
    /// HashSet optimized for storing string or binary values that can produce that
    33
    /// the final set as a `GenericBinaryViewArray` with minimal copies.
    34
    #[derive(Debug)]
    35
    pub struct ArrowBytesViewSet(ArrowBytesViewMap<()>);
    36
    37
    impl ArrowBytesViewSet {
    38
    0
        pub fn new(output_type: OutputType) -> Self {
    39
    0
            Self(ArrowBytesViewMap::new(output_type))
    40
    0
        }
    41
    42
        /// Inserts each value from `values` into the set
    43
    0
        pub fn insert(&mut self, values: &ArrayRef) {
    44
    0
            fn make_payload_fn(_value: Option<&[u8]>) {}
    45
    0
            fn observe_payload_fn(_payload: ()) {}
    46
    0
            self.0
    47
    0
                .insert_if_new(values, make_payload_fn, observe_payload_fn);
    48
    0
        }
    49
    50
        /// Return the contents of this map and replace it with a new empty map with
    51
        /// the same output type
    52
    0
        pub fn take(&mut self) -> Self {
    53
    0
            let mut new_self = Self::new(self.0.output_type);
    54
    0
            std::mem::swap(self, &mut new_self);
    55
    0
            new_self
    56
    0
        }
    57
    58
        /// Converts this set into a `StringViewArray` or `BinaryViewArray`
    59
        /// containing each distinct value that was interned.
    60
        /// This is done without copying the values.
    61
    0
        pub fn into_state(self) -> ArrayRef {
    62
    0
            self.0.into_state()
    63
    0
        }
    64
    65
        /// Returns the total number of distinct values (including nulls) seen so far
    66
    0
        pub fn len(&self) -> usize {
    67
    0
            self.0.len()
    68
    0
        }
    69
    70
    0
        pub fn is_empty(&self) -> bool {
    71
    0
            self.0.is_empty()
    72
    0
        }
    73
    74
        /// returns the total number of distinct values (not including nulls) seen so far
    75
    0
        pub fn non_null_len(&self) -> usize {
    76
    0
            self.0.non_null_len()
    77
    0
        }
    78
    79
        /// Return the total size, in bytes, of memory used to store the data in
    80
        /// this set, not including `self`
    81
    0
        pub fn size(&self) -> usize {
    82
    0
            self.0.size()
    83
    0
        }
    84
    }
    85
    86
    /// Optimized map for storing Arrow "byte view" types (`StringView`, `BinaryView`)
    87
    /// values that can produce the set of keys on
    88
    /// output as `GenericBinaryViewArray` without copies.
    89
    ///
    90
    /// Equivalent to `HashSet<String, V>` but with better performance if you need
    91
    /// to emit the keys as an Arrow `StringViewArray` / `BinaryViewArray`. For other
    92
    /// purposes it is the same as a `HashMap<String, V>`
    93
    ///
    94
    /// # Generic Arguments
    95
    ///
    96
    /// * `V`: payload type
    97
    ///
    98
    /// # Description
    99
    ///
    100
    /// This is a specialized HashMap with the following properties:
    101
    ///
    102
    /// 1. Optimized for storing and emitting Arrow byte types  (e.g.
    103
    ///    `StringViewArray` / `BinaryViewArray`) very efficiently by minimizing copying of
    104
    ///    the string values themselves, both when inserting and when emitting the
    105
    ///    final array.
    106
    ///
    107
    /// 2. Retains the insertion order of entries in the final array. The values are
    108
    ///    in the same order as they were inserted.
    109
    ///
    110
    /// Note this structure can be used as a `HashSet` by specifying the value type
    111
    /// as `()`, as is done by [`ArrowBytesViewSet`].
    112
    ///
    113
    /// This map is used by the special `COUNT DISTINCT` aggregate function to
    114
    /// store the distinct values, and by the `GROUP BY` operator to store
    115
    /// group values when they are a single string array.
    116
    pub struct ArrowBytesViewMap<V>
    117
    where
    118
        V: Debug + PartialEq + Eq + Clone + Copy + Default,
    119
    {
    120
        /// Should the output be StringView or BinaryView?
    121
        output_type: OutputType,
    122
        /// Underlying hash set for each distinct value
    123
        map: hashbrown::hash_table::HashTable<Entry<V>>,
    124
        /// Total size of the map in bytes
    125
        map_size: usize,
    126
    127
        /// Builder for output array
    128
        builder: GenericByteViewBuilder<BinaryViewType>,
    129
        /// random state used to generate hashes
    130
        random_state: RandomState,
    131
        /// buffer that stores hash values (reused across batches to save allocations)
    132
        hashes_buffer: Vec<u64>,
    133
        /// `(payload, null_index)` for the 'null' value, if any
    134
        /// NOTE null_index is the logical index in the final array, not the index
    135
        /// in the buffer
    136
        null: Option<(V, usize)>,
    137
    }
    138
    139
    /// The size, in number of entries, of the initial hash table
    140
    const INITIAL_MAP_CAPACITY: usize = 512;
    141
    142
    impl<V> ArrowBytesViewMap<V>
    143
    where
    144
        V: Debug + PartialEq + Eq + Clone + Copy + Default,
    145
    {
    146
    0
        pub fn new(output_type: OutputType) -> Self {
    147
    0
            Self {
    148
    0
                output_type,
    149
    0
                map: hashbrown::hash_table::HashTable::with_capacity(INITIAL_MAP_CAPACITY),
    150
    0
                map_size: 0,
    151
    0
                builder: GenericByteViewBuilder::new(),
    152
    0
                random_state: RandomState::new(),
    153
    0
                hashes_buffer: vec![],
    154
    0
                null: None,
    155
    0
            }
    156
    0
        }
    157
    158
        /// Return the contents of this map and replace it with a new empty map with
    159
        /// the same output type
    160
        pub fn take(&mut self) -> Self {
    161
            let mut new_self = Self::new(self.output_type);
    162
            std::mem::swap(self, &mut new_self);
    163
            new_self
    164
        }
    165
    166
        /// Inserts each value from `values` into the map, invoking `payload_fn` for
    167
        /// each value if *not* already present, deferring the allocation of the
    168
        /// payload until it is needed.
    169
        ///
    170
        /// Note that this is different than a normal map that would replace the
    171
        /// existing entry
    172
        ///
    173
        /// # Arguments:
    174
        ///
    175
        /// `values`: array whose values are inserted
    176
        ///
    177
        /// `make_payload_fn`:  invoked for each value that is not already present
    178
        /// to create the payload, in order of the values in `values`
    179
        ///
    180
        /// `observe_payload_fn`: invoked once, for each value in `values`, that was
    181
        /// already present in the map, with corresponding payload value.
    182
        ///
    183
        /// # Returns
    184
        ///
    185
        /// The payload value for the entry, either the existing value or
    186
        /// the newly inserted value
    187
        ///
    188
        /// # Safety:
    189
        ///
    190
        /// Note that `make_payload_fn` and `observe_payload_fn` are only invoked
    191
        /// with valid values from `values`, not for the `NULL` value.
    192
    0
        pub fn insert_if_new<MP, OP>(
    193
    0
            &mut self,
    194
    0
            values: &ArrayRef,
    195
    0
            make_payload_fn: MP,
    196
    0
            observe_payload_fn: OP,
    197
    0
        ) where
    198
    0
            MP: FnMut(Option<&[u8]>) -> V,
    199
    0
            OP: FnMut(V),
    200
        {
    201
            // Sanity check array type
    202
    0
            match self.output_type {
    203
                OutputType::BinaryView => {
    204
    0
                    assert!(matches!(values.data_type(), DataType::BinaryView));
    205
    0
                    self.insert_if_new_inner::<MP, OP, BinaryViewType>(
    206
    0
                        values,
    207
    0
                        make_payload_fn,
    208
    0
                        observe_payload_fn,
    209
                    )
    210
                }
    211
                OutputType::Utf8View => {
    212
    0
                    assert!(matches!(values.data_type(), DataType::Utf8View));
    213
    0
                    self.insert_if_new_inner::<MP, OP, StringViewType>(
    214
    0
                        values,
    215
    0
                        make_payload_fn,
    216
    0
                        observe_payload_fn,
    217
                    )
    218
                }
    219
    0
                _ => unreachable!("Utf8/Binary should use `ArrowBytesSet`"),
    220
            };
    221
    0
        }
    222
    223
        /// Generic version of [`Self::insert_if_new`] that handles `ByteViewType`
    224
        /// (both StringView and BinaryView)
    225
        ///
    226
        /// Note this is the only function that is generic on [`ByteViewType`], which
    227
        /// avoids having to template the entire structure,  making the code
    228
        /// simpler and understand and reducing code bloat due to duplication.
    229
        ///
    230
        /// See comments on `insert_if_new` for more details
    231
    0
        fn insert_if_new_inner<MP, OP, B>(
    232
    0
            &mut self,
    233
    0
            values: &ArrayRef,
    234
    0
            mut make_payload_fn: MP,
    235
    0
            mut observe_payload_fn: OP,
    236
    0
        ) where
    237
    0
            MP: FnMut(Option<&[u8]>) -> V,
    238
    0
            OP: FnMut(V),
    239
    0
            B: ByteViewType,
    240
        {
    241
            // step 1: compute hashes
    242
    0
            let batch_hashes = &mut self.hashes_buffer;
    243
    0
            batch_hashes.clear();
    244
    0
            batch_hashes.resize(values.len(), 0);
    245
    0
            create_hashes([values], &self.random_state, batch_hashes)
    246
                // hash is supported for all types and create_hashes only
    247
                // returns errors for unsupported types
    248
    0
                .unwrap();
    249
    250
            // step 2: insert each value into the set, if not already present
    251
    0
            let values = values.as_byte_view::<B>();
    252
    253
            // Ensure lengths are equivalent
    254
    0
            assert_eq!(values.len(), batch_hashes.len());
    255
    256
    0
            for (value, &hash) in values.iter().zip(batch_hashes.iter()) {
    257
                // handle null value
    258
    0
                let Some(value) = value else {
    259
    0
                    let payload = if let Some(&(payload, _offset)) = self.null.as_ref() {
    260
    0
                        payload
    261
                    } else {
    262
    0
                        let payload = make_payload_fn(None);
    263
    0
                        let null_index = self.builder.len();
    264
    0
                        self.builder.append_null();
    265
    0
                        self.null = Some((payload, null_index));
    266
    0
                        payload
    267
                    };
    268
    0
                    observe_payload_fn(payload);
    269
    0
                    continue;
    270
                };
    271
    272
                // get the value as bytes
    273
    0
                let value: &[u8] = value.as_ref();
    274
    275
    0
                let entry = self.map.find_mut(hash, |header| {
    276
    0
                    let v = self.builder.get_value(header.view_idx);
    277
    278
    0
                    if v.len() != value.len() {
    279
    0
                        return false;
    280
    0
                    }
    281
    282
    0
                    v == value
    283
    0
                });
    284
    285
    0
                let payload = if let Some(entry) = entry {
    286
    0
                    entry.payload
    287
                } else {
    288
                    // no existing value, make a new one.
    289
    0
                    let payload = make_payload_fn(Some(value));
    290
    291
    0
                    let inner_view_idx = self.builder.len();
    292
    0
                    let new_header = Entry {
    293
    0
                        view_idx: inner_view_idx,
    294
    0
                        hash,
    295
    0
                        payload,
    296
    0
                    };
    297
    298
    0
                    self.builder.append_value(value);
    299
    300
    0
                    self.map
    301
    0
                        .insert_accounted(new_header, |h| h.hash, &mut self.map_size);
    302
    0
                    payload
    303
                };
    304
    0
                observe_payload_fn(payload);
    305
            }
    306
    0
        }
    307
    308
        /// Converts this set into a `StringViewArray`, or `BinaryViewArray`,
    309
        /// containing each distinct value
    310
        /// that was inserted. This is done without copying the values.
    311
        ///
    312
        /// The values are guaranteed to be returned in the same order in which
    313
        /// they were first seen.
    314
    0
        pub fn into_state(self) -> ArrayRef {
    315
    0
            let mut builder = self.builder;
    316
    0
            match self.output_type {
    317
                OutputType::BinaryView => {
    318
    0
                    let array = builder.finish();
    319
    320
    0
                    Arc::new(array)
    321
                }
    322
                OutputType::Utf8View => {
    323
                    // SAFETY:
    324
                    // we asserted the input arrays were all the correct type and
    325
                    // thus since all the values that went in were valid (e.g. utf8)
    326
                    // so are all the values that come out
    327
    0
                    let array = builder.finish();
    328
    0
                    let array = unsafe { array.to_string_view_unchecked() };
    329
    0
                    Arc::new(array)
    330
                }
    331
                _ => {
    332
    0
                    unreachable!("Utf8/Binary should use `ArrowBytesMap`")
    333
                }
    334
            }
    335
    0
        }
    336
    337
        /// Total number of entries (including null, if present)
    338
    0
        pub fn len(&self) -> usize {
    339
    0
            self.non_null_len() + self.null.map(|_| 1).unwrap_or(0)
    340
    0
        }
    341
    342
        /// Is the set empty?
    343
    0
        pub fn is_empty(&self) -> bool {
    344
    0
            self.map.is_empty() && self.null.is_none()
    345
    0
        }
    346
    347
        /// Number of non null entries
    348
    0
        pub fn non_null_len(&self) -> usize {
    349
    0
            self.map.len()
    350
    0
        }
    351
    352
        /// Return the total size, in bytes, of memory used to store the data in
    353
        /// this set, not including `self`
    354
    0
        pub fn size(&self) -> usize {
    355
    0
            self.map_size
    356
    0
                + self.builder.allocated_size()
    357
    0
                + self.hashes_buffer.allocated_size()
    358
    0
        }
    359
    }
    360
    361
    impl<V> Debug for ArrowBytesViewMap<V>
    362
    where
    363
        V: Debug + PartialEq + Eq + Clone + Copy + Default,
    364
    {
    365
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
    366
            f.debug_struct("ArrowBytesMap")
    367
                .field("map", &"<map>")
    368
                .field("map_size", &self.map_size)
    369
                .field("view_builder", &self.builder)
    370
                .field("random_state", &self.random_state)
    371
                .field("hashes_buffer", &self.hashes_buffer)
    372
                .finish()
    373
        }
    374
    }
    375
    376
    /// Entry in the hash table -- see [`ArrowBytesViewMap`] for more details
    377
    #[derive(Debug, PartialEq, Eq, Hash, Clone, Copy)]
    378
    struct Entry<V>
    379
    where
    380
        V: Debug + PartialEq + Eq + Clone + Copy + Default,
    381
    {
    382
        /// The idx into the views array
    383
        view_idx: usize,
    384
    385
        hash: u64,
    386
    387
        /// value stored by the entry
    388
        payload: V,
    389
    }
    390
    391
    #[cfg(test)]
    392
    mod tests {
    393
        use arrow::array::{BinaryViewArray, GenericByteViewArray, StringViewArray};
    394
        use datafusion_common::HashMap;
    395
    396
        use super::*;
    397
    398
        // asserts that the set contains the expected strings, in the same order
    399
        fn assert_set(set: ArrowBytesViewSet, expected: &[Option<&str>]) {
    400
            let strings = set.into_state();
    401
            let strings = strings.as_string_view();
    402
            let state = strings.into_iter().collect::<Vec<_>>();
    403
            assert_eq!(state, expected);
    404
        }
    405
    406
        #[test]
    407
        fn string_view_set_empty() {
    408
            let mut set = ArrowBytesViewSet::new(OutputType::Utf8View);
    409
            let array: ArrayRef = Arc::new(StringViewArray::new_null(0));
    410
            set.insert(&array);
    411
            assert_eq!(set.len(), 0);
    412
            assert_eq!(set.non_null_len(), 0);
    413
            assert_set(set, &[]);
    414
        }
    415
    416
        #[test]
    417
        fn string_view_set_one_null() {
    418
            let mut set = ArrowBytesViewSet::new(OutputType::Utf8View);
    419
            let array: ArrayRef = Arc::new(StringViewArray::new_null(1));
    420
            set.insert(&array);
    421
            assert_eq!(set.len(), 1);
    422
            assert_eq!(set.non_null_len(), 0);
    423
            assert_set(set, &[None]);
    424
        }
    425
    426
        #[test]
    427
        fn string_view_set_many_null() {
    428
            let mut set = ArrowBytesViewSet::new(OutputType::Utf8View);
    429
            let array: ArrayRef = Arc::new(StringViewArray::new_null(11));
    430
            set.insert(&array);
    431
            assert_eq!(set.len(), 1);
    432
            assert_eq!(set.non_null_len(), 0);
    433
            assert_set(set, &[None]);
    434
        }
    435
    436
        #[test]
    437
        fn test_string_view_set_basic() {
    438
            // basic test for mixed small and large string values
    439
            let values = GenericByteViewArray::from(vec![
    440
                Some("a"),
    441
                Some("b"),
    442
                Some("CXCCCCCCCCAABB"), // 14 bytes
    443
                Some(""),
    444
                Some("cbcxx"), // 5 bytes
    445
                None,
    446
                Some("AAAAAAAA"),     // 8 bytes
    447
                Some("BBBBBQBBBAAA"), // 12 bytes
    448
                Some("a"),
    449
                Some("cbcxx"),
    450
                Some("b"),
    451
                Some("cbcxx"),
    452
                Some(""),
    453
                None,
    454
                Some("BBBBBQBBBAAA"),
    455
                Some("BBBBBQBBBAAA"),
    456
                Some("AAAAAAAA"),
    457
                Some("CXCCCCCCCCAABB"),
    458
            ]);
    459
    460
            let mut set = ArrowBytesViewSet::new(OutputType::Utf8View);
    461
            let array: ArrayRef = Arc::new(values);
    462
            set.insert(&array);
    463
            // values mut appear be in the order they were inserted
    464
            assert_set(
    465
                set,
    466
                &[
    467
                    Some("a"),
    468
                    Some("b"),
    469
                    Some("CXCCCCCCCCAABB"),
    470
                    Some(""),
    471
                    Some("cbcxx"),
    472
                    None,
    473
                    Some("AAAAAAAA"),
    474
                    Some("BBBBBQBBBAAA"),
    475
                ],
    476
            );
    477
        }
    478
    479
        #[test]
    480
        fn test_string_set_non_utf8() {
    481
            // basic test for mixed small and large string values
    482
            let values = GenericByteViewArray::from(vec![
    483
                Some("a"),
    484
                Some("✨🔥✨🔥✨🔥✨🔥✨🔥✨🔥✨🔥✨🔥"),
    485
                Some("🔥"),
    486
                Some("✨✨✨"),
    487
                Some("foobarbaz"),
    488
                Some("🔥"),
    489
                Some("✨🔥✨🔥✨🔥✨🔥✨🔥✨🔥✨🔥✨🔥"),
    490
            ]);
    491
    492
            let mut set = ArrowBytesViewSet::new(OutputType::Utf8View);
    493
            let array: ArrayRef = Arc::new(values);
    494
            set.insert(&array);
    495
            // strings mut appear be in the order they were inserted
    496
            assert_set(
    497
                set,
    498
                &[
    499
                    Some("a"),
    500
                    Some("✨🔥✨🔥✨🔥✨🔥✨🔥✨🔥✨🔥✨🔥"),
    501
                    Some("🔥"),
    502
                    Some("✨✨✨"),
    503
                    Some("foobarbaz"),
    504
                ],
    505
            );
    506
        }
    507
    508
        // Test use of binary output type
    509
        #[test]
    510
        fn test_binary_set() {
    511
            let v: Vec<Option<&[u8]>> = vec![
    512
                Some(b"a"),
    513
                Some(b"CXCCCCCCCCCCCCC"),
    514
                None,
    515
                Some(b"CXCCCCCCCCCCCCC"),
    516
            ];
    517
            let values: ArrayRef = Arc::new(BinaryViewArray::from(v));
    518
    519
            let expected: Vec<Option<&[u8]>> =
    520
                vec![Some(b"a"), Some(b"CXCCCCCCCCCCCCC"), None];
    521
            let expected: ArrayRef = Arc::new(GenericByteViewArray::from(expected));
    522
    523
            let mut set = ArrowBytesViewSet::new(OutputType::BinaryView);
    524
            set.insert(&values);
    525
            assert_eq!(&set.into_state(), &expected);
    526
        }
    527
    528
        // inserting strings into the set does not increase reported memory
    529
        #[test]
    530
        fn test_string_set_memory_usage() {
    531
            let strings1 = StringViewArray::from(vec![
    532
                Some("a"),
    533
                Some("b"),
    534
                Some("CXCCCCCCCCCCC"), // 13 bytes
    535
                Some("AAAAAAAA"),      // 8 bytes
    536
                Some("BBBBBQBBB"),     // 9 bytes
    537
            ]);
    538
            let total_strings1_len = strings1
    539
                .iter()
    540
                .map(|s| s.map(|s| s.len()).unwrap_or(0))
    541
                .sum::<usize>();
    542
            let values1: ArrayRef = Arc::new(StringViewArray::from(strings1));
    543
    544
            // Much larger strings in strings2
    545
            let strings2 = StringViewArray::from(vec![
    546
                "FOO".repeat(1000),
    547
                "BAR larger than 12 bytes.".repeat(100_000),
    548
                "more unique.".repeat(1000),
    549
                "more unique2.".repeat(1000),
    550
                "FOO".repeat(3000),
    551
            ]);
    552
            let total_strings2_len = strings2
    553
                .iter()
    554
                .map(|s| s.map(|s| s.len()).unwrap_or(0))
    555
                .sum::<usize>();
    556
            let values2: ArrayRef = Arc::new(StringViewArray::from(strings2));
    557
    558
            let mut set = ArrowBytesViewSet::new(OutputType::Utf8View);
    559
            let size_empty = set.size();
    560
    561
            set.insert(&values1);
    562
            let size_after_values1 = set.size();
    563
            assert!(size_empty < size_after_values1);
    564
            assert!(
    565
                size_after_values1 > total_strings1_len,
    566
                "expect {size_after_values1} to be more than {total_strings1_len}"
    567
            );
    568
            assert!(size_after_values1 < total_strings1_len + total_strings2_len);
    569
    570
            // inserting the same strings should not affect the size
    571
            set.insert(&values1);
    572
            assert_eq!(set.size(), size_after_values1);
    573
            assert_eq!(set.len(), 5);
    574
    575
            // inserting the large strings should increase the reported size
    576
            set.insert(&values2);
    577
            let size_after_values2 = set.size();
    578
            assert!(size_after_values2 > size_after_values1);
    579
    580
            assert_eq!(set.len(), 10);
    581
        }
    582
    583
        #[derive(Debug, PartialEq, Eq, Default, Clone, Copy)]
    584
        struct TestPayload {
    585
            // store the string value to check against input
    586
            index: usize, // store the index of the string (each new string gets the next sequential input)
    587
        }
    588
    589
        /// Wraps an [`ArrowBytesViewMap`], validating its invariants
    590
        struct TestMap {
    591
            map: ArrowBytesViewMap<TestPayload>,
    592
            // stores distinct strings seen, in order
    593
            strings: Vec<Option<String>>,
    594
            // map strings to index in strings
    595
            indexes: HashMap<Option<String>, usize>,
    596
        }
    597
    598
        impl TestMap {
    599
            /// creates a map with TestPayloads for the given strings and then
    600
            /// validates the payloads
    601
            fn new() -> Self {
    602
                Self {
    603
                    map: ArrowBytesViewMap::new(OutputType::Utf8View),
    604
                    strings: vec![],
    605
                    indexes: HashMap::new(),
    606
                }
    607
            }
    608
    609
            /// Inserts strings into the map
    610
            fn insert(&mut self, strings: &[Option<&str>]) {
    611
                let string_array = StringViewArray::from(strings.to_vec());
    612
                let arr: ArrayRef = Arc::new(string_array);
    613
    614
                let mut next_index = self.indexes.len();
    615
                let mut actual_new_strings = vec![];
    616
                let mut actual_seen_indexes = vec![];
    617
                // update self with new values, keeping track of newly added values
    618
                for str in strings {
    619
                    let str = str.map(|s| s.to_string());
    620
                    let index = self.indexes.get(&str).cloned().unwrap_or_else(|| {
    621
                        actual_new_strings.push(str.clone());
    622
                        let index = self.strings.len();
    623
                        self.strings.push(str.clone());
    624
                        self.indexes.insert(str, index);
    625
                        index
    626
                    });
    627
                    actual_seen_indexes.push(index);
    628
                }
    629
    630
                // insert the values into the map, recording what we did
    631
                let mut seen_new_strings = vec![];
    632
                let mut seen_indexes = vec![];
    633
                self.map.insert_if_new(
    634
                    &arr,
    635
                    |s| {
    636
                        let value = s
    637
                            .map(|s| String::from_utf8(s.to_vec()).expect("Non utf8 string"));
    638
                        let index = next_index;
    639
                        next_index += 1;
    640
                        seen_new_strings.push(value);
    641
                        TestPayload { index }
    642
                    },
    643
                    |payload| {
    644
                        seen_indexes.push(payload.index);
    645
                    },
    646
                );
    647
    648
                assert_eq!(actual_seen_indexes, seen_indexes);
    649
                assert_eq!(actual_new_strings, seen_new_strings);
    650
            }
    651
    652
            /// Call `self.map.into_array()` validating that the strings are in the same
    653
            /// order as they were inserted
    654
            fn into_array(self) -> ArrayRef {
    655
                let Self {
    656
                    map,
    657
                    strings,
    658
                    indexes: _,
    659
                } = self;
    660
    661
                let arr = map.into_state();
    662
                let expected: ArrayRef = Arc::new(StringViewArray::from(strings));
    663
                assert_eq!(&arr, &expected);
    664
                arr
    665
            }
    666
        }
    667
    668
        #[test]
    669
        fn test_map() {
    670
            let input = vec![
    671
                // Note mix of short/long strings
    672
                Some("A"),
    673
                Some("bcdefghijklmnop1234567"),
    674
                Some("X"),
    675
                Some("Y"),
    676
                None,
    677
                Some("qrstuvqxyzhjwya"),
    678
                Some("✨🔥"),
    679
                Some("🔥"),
    680
                Some("🔥🔥🔥🔥🔥🔥"),
    681
            ];
    682
    683
            let mut test_map = TestMap::new();
    684
            test_map.insert(&input);
    685
            test_map.insert(&input); // put it in twice
    686
            let expected_output: ArrayRef = Arc::new(StringViewArray::from(input));
    687
            assert_eq!(&test_map.into_array(), &expected_output);
    688
        }
    689
    }