## 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/common/src/utils/mod.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
    //! This module provides the bisect function, which implements binary search.
    19
    20
    pub mod expr;
    21
    pub mod memory;
    22
    pub mod proxy;
    23
    pub mod string_utils;
    24
    25
    use crate::assert_or_internal_err;
    26
    use crate::error::{_exec_datafusion_err, _internal_datafusion_err};
    27
    use crate::{DataFusionError, Result, ScalarValue};
    28
    use arrow::array::{
    29
        cast::AsArray, Array, ArrayRef, FixedSizeListArray, LargeListArray, ListArray,
    30
        OffsetSizeTrait,
    31
    };
    32
    use arrow::buffer::OffsetBuffer;
    33
    use arrow::compute::{partition, SortColumn, SortOptions};
    34
    use arrow::datatypes::{DataType, Field, SchemaRef};
    35
    #[cfg(feature = "sql")]
    36
    use sqlparser::{ast::Ident, dialect::GenericDialect, parser::Parser};
    37
    use std::borrow::{Borrow, Cow};
    38
    use std::cmp::{min, Ordering};
    39
    use std::collections::HashSet;
    40
    use std::num::NonZero;
    41
    use std::ops::Range;
    42
    use std::sync::Arc;
    43
    use std::thread::available_parallelism;
    44
    45
    /// Applies an optional projection to a [`SchemaRef`], returning the
    46
    /// projected schema
    47
    ///
    48
    /// Example:
    49
    /// ```
    50
    /// use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
    51
    /// use datafusion_common::project_schema;
    52
    ///
    53
    /// // Schema with columns 'a', 'b', and 'c'
    54
    /// let schema = SchemaRef::new(Schema::new(vec![
    55
    ///     Field::new("a", DataType::Int32, true),
    56
    ///     Field::new("b", DataType::Int64, true),
    57
    ///     Field::new("c", DataType::Utf8, true),
    58
    /// ]));
    59
    ///
    60
    /// // Pick columns 'c' and 'b'
    61
    /// let projection = Some(vec![2, 1]);
    62
    /// let projected_schema = project_schema(&schema, projection.as_ref()).unwrap();
    63
    ///
    64
    /// let expected_schema = SchemaRef::new(Schema::new(vec![
    65
    ///     Field::new("c", DataType::Utf8, true),
    66
    ///     Field::new("b", DataType::Int64, true),
    67
    /// ]));
    68
    ///
    69
    /// assert_eq!(projected_schema, expected_schema);
    70
    /// ```
    71
    0
    pub fn project_schema(
    72
    0
        schema: &SchemaRef,
    73
    0
        projection: Option<&Vec<usize>>,
    74
    0
    ) -> Result<SchemaRef> {
    75
    0
        let schema = match projection {
    76
    0
            Some(columns) => Arc::new(schema.project(columns)?),
    77
    0
            None => Arc::clone(schema),
    78
        };
    79
    0
        Ok(schema)
    80
    0
    }
    81
    82
    /// Extracts a row at the specified index from a set of columns and stores it in the provided buffer.
    83
    0
    pub fn extract_row_at_idx_to_buf(
    84
    0
        columns: &[ArrayRef],
    85
    0
        idx: usize,
    86
    0
        buf: &mut Vec<ScalarValue>,
    87
    0
    ) -> Result<()> {
    88
    0
        buf.clear();
    89
    90
    0
        let iter = columns
    91
    0
            .iter()
    92
    0
            .map(|arr| ScalarValue::try_from_array(arr, idx));
    93
    0
        for v in iter.into_iter() {
    94
    0
            buf.push(v?);
    95
        }
    96
    97
    0
        Ok(())
    98
    0
    }
    99
    /// Given column vectors, returns row at `idx`.
    100
    0
    pub fn get_row_at_idx(columns: &[ArrayRef], idx: usize) -> Result<Vec<ScalarValue>> {
    101
    0
        columns
    102
    0
            .iter()
    103
    0
            .map(|arr| ScalarValue::try_from_array(arr, idx))
    104
    0
            .collect()
    105
    0
    }
    106
    107
    /// This function compares two tuples depending on the given sort options.
    108
    0
    pub fn compare_rows(
    109
    0
        x: &[ScalarValue],
    110
    0
        y: &[ScalarValue],
    111
    0
        sort_options: &[SortOptions],
    112
    0
    ) -> Result<Ordering> {
    113
    0
        let zip_it = x.iter().zip(y.iter()).zip(sort_options.iter());
    114
        // Preserving lexical ordering.
    115
    0
        for ((lhs, rhs), sort_options) in zip_it {
    116
            // Consider all combinations of NULLS FIRST/LAST and ASC/DESC configurations.
    117
    0
            let result = match (lhs.is_null(), rhs.is_null(), sort_options.nulls_first) {
    118
    0
                (true, false, false) | (false, true, true) => Ordering::Greater,
    119
    0
                (true, false, true) | (false, true, false) => Ordering::Less,
    120
                (false, false, _) => {
    121
    0
                    if sort_options.descending {
    122
    0
                        rhs.try_cmp(lhs)?
    123
                    } else {
    124
    0
                        lhs.try_cmp(rhs)?
    125
                    }
    126
                }
    127
    0
                (true, true, _) => continue,
    128
            };
    129
    0
            if result != Ordering::Equal {
    130
    0
                return Ok(result);
    131
    0
            }
    132
        }
    133
    0
        Ok(Ordering::Equal)
    134
    0
    }
    135
    136
    /// This function searches for a tuple of given values (`target`) among the given
    137
    /// rows (`item_columns`) using the bisection algorithm. It assumes that `item_columns`
    138
    /// is sorted according to `sort_options` and returns the insertion index of `target`.
    139
    /// Template argument `SIDE` being `true`/`false` means left/right insertion.
    140
    0
    pub fn bisect<const SIDE: bool>(
    141
    0
        item_columns: &[ArrayRef],
    142
    0
        target: &[ScalarValue],
    143
    0
        sort_options: &[SortOptions],
    144
    0
    ) -> Result<usize> {
    145
    0
        let low: usize = 0;
    146
    0
        let high: usize = item_columns
    147
    0
            .first()
    148
    0
            .ok_or_else(|| _internal_datafusion_err!("Column array shouldn't be empty"))?
    149
    0
            .len();
    150
    0
        let compare_fn = |current: &[ScalarValue], target: &[ScalarValue]| {
    151
    0
            let cmp = compare_rows(current, target, sort_options)?;
    152
    0
            Ok(if SIDE { cmp.is_lt() } else { cmp.is_le() })
    153
    0
        };
    154
    0
        find_bisect_point(item_columns, target, compare_fn, low, high)
    155
    0
    }
    156
    157
    /// This function searches for a tuple of given values (`target`) among a slice of
    158
    /// the given rows (`item_columns`) using the bisection algorithm. The slice starts
    159
    /// at the index `low` and ends at the index `high`. The boolean-valued function
    160
    /// `compare_fn` specifies whether we bisect on the left (by returning `false`),
    161
    /// or on the right (by returning `true`) when we compare the target value with
    162
    /// the current value as we iteratively bisect the input.
    163
    0
    pub fn find_bisect_point<F>(
    164
    0
        item_columns: &[ArrayRef],
    165
    0
        target: &[ScalarValue],
    166
    0
        compare_fn: F,
    167
    0
        mut low: usize,
    168
    0
        mut high: usize,
    169
    0
    ) -> Result<usize>
    170
    0
    where
    171
    0
        F: Fn(&[ScalarValue], &[ScalarValue]) -> Result<bool>,
    172
    {
    173
    0
        while low < high {
    174
    0
            let mid = ((high - low) / 2) + low;
    175
    0
            let val = get_row_at_idx(item_columns, mid)?;
    176
    0
            if compare_fn(&val, target)? {
    177
    0
                low = mid + 1;
    178
    0
            } else {
    179
    0
                high = mid;
    180
    0
            }
    181
        }
    182
    0
        Ok(low)
    183
    0
    }
    184
    185
    /// This function searches for a tuple of given values (`target`) among the given
    186
    /// rows (`item_columns`) via a linear scan. It assumes that `item_columns` is sorted
    187
    /// according to `sort_options` and returns the insertion index of `target`.
    188
    /// Template argument `SIDE` being `true`/`false` means left/right insertion.
    189
    0
    pub fn linear_search<const SIDE: bool>(
    190
    0
        item_columns: &[ArrayRef],
    191
    0
        target: &[ScalarValue],
    192
    0
        sort_options: &[SortOptions],
    193
    0
    ) -> Result<usize> {
    194
    0
        let low: usize = 0;
    195
    0
        let high: usize = item_columns
    196
    0
            .first()
    197
    0
            .ok_or_else(|| _internal_datafusion_err!("Column array shouldn't be empty"))?
    198
    0
            .len();
    199
    0
        let compare_fn = |current: &[ScalarValue], target: &[ScalarValue]| {
    200
    0
            let cmp = compare_rows(current, target, sort_options)?;
    201
    0
            Ok(if SIDE { cmp.is_lt() } else { cmp.is_le() })
    202
    0
        };
    203
    0
        search_in_slice(item_columns, target, compare_fn, low, high)
    204
    0
    }
    205
    206
    /// This function searches for a tuple of given values (`target`) among a slice of
    207
    /// the given rows (`item_columns`) via a linear scan. The slice starts at the index
    208
    /// `low` and ends at the index `high`. The boolean-valued function `compare_fn`
    209
    /// specifies the stopping criterion.
    210
    0
    pub fn search_in_slice<F>(
    211
    0
        item_columns: &[ArrayRef],
    212
    0
        target: &[ScalarValue],
    213
    0
        compare_fn: F,
    214
    0
        mut low: usize,
    215
    0
        high: usize,
    216
    0
    ) -> Result<usize>
    217
    0
    where
    218
    0
        F: Fn(&[ScalarValue], &[ScalarValue]) -> Result<bool>,
    219
    {
    220
    0
        while low < high {
    221
    0
            let val = get_row_at_idx(item_columns, low)?;
    222
    0
            if !compare_fn(&val, target)? {
    223
    0
                break;
    224
    0
            }
    225
    0
            low += 1;
    226
        }
    227
    0
        Ok(low)
    228
    0
    }
    229
    230
    /// Given a list of 0 or more already sorted columns, finds the
    231
    /// partition ranges that would partition equally across columns.
    232
    ///
    233
    /// See [`partition`] for more details.
    234
    0
    pub fn evaluate_partition_ranges(
    235
    0
        num_rows: usize,
    236
    0
        partition_columns: &[SortColumn],
    237
    0
    ) -> Result<Vec<Range<usize>>> {
    238
    0
        Ok(if partition_columns.is_empty() {
    239
    0
            vec![Range {
    240
    0
                start: 0,
    241
    0
                end: num_rows,
    242
    0
            }]
    243
        } else {
    244
    0
            let cols: Vec<_> = partition_columns
    245
    0
                .iter()
    246
    0
                .map(|x| Arc::clone(&x.values))
    247
    0
                .collect();
    248
    0
            partition(&cols)?.ranges()
    249
        })
    250
    0
    }
    251
    252
    /// Wraps identifier string in double quotes, escaping any double quotes in
    253
    /// the identifier by replacing it with two double quotes
    254
    ///
    255
    /// e.g. identifier `tab.le"name` becomes `"tab.le""name"`
    256
    0
    pub fn quote_identifier(s: &str) -> Cow<'_, str> {
    257
    0
        if needs_quotes(s) {
    258
    0
            Cow::Owned(format!("\"{}\"", s.replace('"', "\"\"")))
    259
        } else {
    260
    0
            Cow::Borrowed(s)
    261
        }
    262
    0
    }
    263
    264
    /// returns true if this identifier needs quotes
    265
    0
    fn needs_quotes(s: &str) -> bool {
    266
    0
        let mut chars = s.chars();
    267
    268
        // first char can not be a number unless escaped
    269
    0
        if let Some(first_char) = chars.next() {
    270
    0
            if !(first_char.is_ascii_lowercase() || first_char == '_') {
    271
    0
                return true;
    272
    0
            }
    273
    0
        }
    274
    275
    0
        !chars.all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_')
    276
    0
    }
    277
    278
    #[cfg(feature = "sql")]
    279
    pub(crate) fn parse_identifiers(s: &str) -> Result<Vec<Ident>> {
    280
        let dialect = GenericDialect;
    281
        let mut parser = Parser::new(&dialect).try_with_sql(s)?;
    282
        let idents = parser.parse_multipart_identifier()?;
    283
        Ok(idents)
    284
    }
    285
    286
    /// Parse a string into a vector of identifiers.
    287
    ///
    288
    /// Note: If ignore_case is false, the string will be normalized to lowercase.
    289
    #[cfg(feature = "sql")]
    290
    pub(crate) fn parse_identifiers_normalized(s: &str, ignore_case: bool) -> Vec<String> {
    291
        parse_identifiers(s)
    292
            .unwrap_or_default()
    293
            .into_iter()
    294
            .map(|id| match id.quote_style {
    295
                Some(_) => id.value,
    296
                None if ignore_case => id.value,
    297
                _ => id.value.to_ascii_lowercase(),
    298
            })
    299
            .collect::<Vec<_>>()
    300
    }
    301
    302
    #[cfg(not(feature = "sql"))]
    303
    0
    pub(crate) fn parse_identifiers(s: &str) -> Result<Vec<String>> {
    304
    0
        let mut result = Vec::new();
    305
    0
        let mut current = String::new();
    306
    0
        let mut in_quotes = false;
    307
    308
    0
        for ch in s.chars() {
    309
    0
            match ch {
    310
    0
                '"' => {
    311
    0
                    in_quotes = !in_quotes;
    312
    0
                    current.push(ch);
    313
    0
                }
    314
    0
                '.' if !in_quotes => {
    315
    0
                    result.push(current.clone());
    316
    0
                    current.clear();
    317
    0
                }
    318
    0
                _ => {
    319
    0
                    current.push(ch);
    320
    0
                }
    321
            }
    322
        }
    323
    324
        // Push the last part if it's not empty
    325
    0
        if !current.is_empty() {
    326
    0
            result.push(current);
    327
    0
        }
    328
    329
    0
        Ok(result)
    330
    0
    }
    331
    332
    #[cfg(not(feature = "sql"))]
    333
    0
    pub(crate) fn parse_identifiers_normalized(s: &str, ignore_case: bool) -> Vec<String> {
    334
    0
        parse_identifiers(s)
    335
    0
            .unwrap_or_default()
    336
    0
            .into_iter()
    337
    0
            .map(|id| {
    338
    0
                let is_double_quoted = if id.len() > 2 {
    339
    0
                    let mut chars = id.chars();
    340
    0
                    chars.next() == Some('"') && chars.last() == Some('"')
    341
                } else {
    342
    0
                    false
    343
                };
    344
    0
                if is_double_quoted {
    345
    0
                    id[1..id.len() - 1].to_string().replace("\"\"", "\"")
    346
    0
                } else if ignore_case {
    347
    0
                    id
    348
                } else {
    349
    0
                    id.to_ascii_lowercase()
    350
                }
    351
    0
            })
    352
    0
            .collect::<Vec<_>>()
    353
    0
    }
    354
    355
    /// This function "takes" the elements at `indices` from the slice `items`.
    356
    0
    pub fn get_at_indices<T: Clone, I: Borrow<usize>>(
    357
    0
        items: &[T],
    358
    0
        indices: impl IntoIterator<Item = I>,
    359
    0
    ) -> Result<Vec<T>> {
    360
    0
        indices
    361
    0
            .into_iter()
    362
    0
            .map(|idx| items.get(*idx.borrow()).cloned())
    363
    0
            .collect::<Option<Vec<T>>>()
    364
    0
            .ok_or_else(|| {
    365
    0
                _exec_datafusion_err!("Expects indices to be in the range of searched vector")
    366
    0
            })
    367
    0
    }
    368
    369
    /// This function finds the longest prefix of the form 0, 1, 2, ... within the
    370
    /// collection `sequence`. Examples:
    371
    /// - For 0, 1, 2, 4, 5; we would produce 3, meaning 0, 1, 2 is the longest satisfying
    372
    ///   prefix.
    373
    /// - For 1, 2, 3, 4; we would produce 0, meaning there is no such prefix.
    374
    0
    pub fn longest_consecutive_prefix<T: Borrow<usize>>(
    375
    0
        sequence: impl IntoIterator<Item = T>,
    376
    0
    ) -> usize {
    377
    0
        let mut count = 0;
    378
    0
        for item in sequence {
    379
    0
            if !count.eq(item.borrow()) {
    380
    0
                break;
    381
    0
            }
    382
    0
            count += 1;
    383
        }
    384
    0
        count
    385
    0
    }
    386
    387
    /// Creates single element [`ListArray`], [`LargeListArray`] and
    388
    /// [`FixedSizeListArray`] from other arrays
    389
    ///
    390
    /// For example this builder can convert `[1, 2, 3]` into `[[1, 2, 3]]`
    391
    ///
    392
    /// # Example
    393
    /// ```
    394
    /// # use std::sync::Arc;
    395
    /// # use arrow::array::{Array, ListArray};
    396
    /// # use arrow::array::types::Int64Type;
    397
    /// # use datafusion_common::utils::SingleRowListArrayBuilder;
    398
    /// // Array is [1, 2, 3]
    399
    /// let arr = ListArray::from_iter_primitive::<Int64Type, _, _>(vec![Some(vec![
    400
    ///     Some(1),
    401
    ///     Some(2),
    402
    ///     Some(3),
    403
    /// ])]);
    404
    /// // Wrap as a list array: [[1, 2, 3]]
    405
    /// let list_arr = SingleRowListArrayBuilder::new(Arc::new(arr)).build_list_array();
    406
    /// assert_eq!(list_arr.len(), 1);
    407
    /// ```
    408
    #[derive(Debug, Clone)]
    409
    pub struct SingleRowListArrayBuilder {
    410
        /// array to be wrapped
    411
        arr: ArrayRef,
    412
        /// Should the resulting array be nullable? Defaults to `true`.
    413
        nullable: bool,
    414
        /// Specify the field name for the resulting array. Defaults to value used in
    415
        /// [`Field::new_list_field`]
    416
        field_name: Option<String>,
    417
    }
    418
    419
    impl SingleRowListArrayBuilder {
    420
        /// Create a new instance of [`SingleRowListArrayBuilder`]
    421
    0
        pub fn new(arr: ArrayRef) -> Self {
    422
    0
            Self {
    423
    0
                arr,
    424
    0
                nullable: true,
    425
    0
                field_name: None,
    426
    0
            }
    427
    0
        }
    428
    429
        /// Set the nullable flag
    430
    0
        pub fn with_nullable(mut self, nullable: bool) -> Self {
    431
    0
            self.nullable = nullable;
    432
    0
            self
    433
    0
        }
    434
    435
        /// sets the field name for the resulting array
    436
    0
        pub fn with_field_name(mut self, field_name: Option<String>) -> Self {
    437
    0
            self.field_name = field_name;
    438
    0
            self
    439
    0
        }
    440
    441
        /// Copies field name and nullable from the specified field
    442
    0
        pub fn with_field(self, field: &Field) -> Self {
    443
    0
            self.with_field_name(Some(field.name().to_owned()))
    444
    0
                .with_nullable(field.is_nullable())
    445
    0
        }
    446
    447
        /// Build a single element [`ListArray`]
    448
    0
        pub fn build_list_array(self) -> ListArray {
    449
    0
            let (field, arr) = self.into_field_and_arr();
    450
    0
            let offsets = OffsetBuffer::from_lengths([arr.len()]);
    451
    0
            ListArray::new(field, offsets, arr, None)
    452
    0
        }
    453
    454
        /// Build a single element [`ListArray`] and wrap as [`ScalarValue::List`]
    455
    0
        pub fn build_list_scalar(self) -> ScalarValue {
    456
    0
            ScalarValue::List(Arc::new(self.build_list_array()))
    457
    0
        }
    458
    459
        /// Build a single element [`LargeListArray`]
    460
    0
        pub fn build_large_list_array(self) -> LargeListArray {
    461
    0
            let (field, arr) = self.into_field_and_arr();
    462
    0
            let offsets = OffsetBuffer::from_lengths([arr.len()]);
    463
    0
            LargeListArray::new(field, offsets, arr, None)
    464
    0
        }
    465
    466
        /// Build a single element [`LargeListArray`] and wrap as [`ScalarValue::LargeList`]
    467
    0
        pub fn build_large_list_scalar(self) -> ScalarValue {
    468
    0
            ScalarValue::LargeList(Arc::new(self.build_large_list_array()))
    469
    0
        }
    470
    471
        /// Build a single element [`FixedSizeListArray`]
    472
    0
        pub fn build_fixed_size_list_array(self, list_size: usize) -> FixedSizeListArray {
    473
    0
            let (field, arr) = self.into_field_and_arr();
    474
    0
            FixedSizeListArray::new(field, list_size as i32, arr, None)
    475
    0
        }
    476
    477
        /// Build a single element [`FixedSizeListArray`] and wrap as [`ScalarValue::FixedSizeList`]
    478
    0
        pub fn build_fixed_size_list_scalar(self, list_size: usize) -> ScalarValue {
    479
    0
            ScalarValue::FixedSizeList(Arc::new(self.build_fixed_size_list_array(list_size)))
    480
    0
        }
    481
    482
        /// Helper function: convert this builder into a tuple of field and array
    483
    0
        fn into_field_and_arr(self) -> (Arc<Field>, ArrayRef) {
    484
            let Self {
    485
    0
                arr,
    486
    0
                nullable,
    487
    0
                field_name,
    488
    0
            } = self;
    489
    0
            let data_type = arr.data_type().to_owned();
    490
    0
            let field = match field_name {
    491
    0
                Some(name) => Field::new(name, data_type, nullable),
    492
    0
                None => Field::new_list_field(data_type, nullable),
    493
            };
    494
    0
            (Arc::new(field), arr)
    495
    0
        }
    496
    }
    497
    498
    /// Wrap arrays into a single element `ListArray`.
    499
    ///
    500
    /// Example:
    501
    /// ```
    502
    /// use arrow::array::{Int32Array, ListArray, ArrayRef};
    503
    /// use arrow::datatypes::{Int32Type, Field};
    504
    /// use std::sync::Arc;
    505
    ///
    506
    /// let arr1 = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef;
    507
    /// let arr2 = Arc::new(Int32Array::from(vec![4, 5, 6])) as ArrayRef;
    508
    ///
    509
    /// let list_arr = datafusion_common::utils::arrays_into_list_array([arr1, arr2]).unwrap();
    510
    ///
    511
    /// let expected = ListArray::from_iter_primitive::<Int32Type, _, _>(
    512
    ///    vec![
    513
    ///     Some(vec![Some(1), Some(2), Some(3)]),
    514
    ///     Some(vec![Some(4), Some(5), Some(6)]),
    515
    ///    ]
    516
    /// );
    517
    ///
    518
    /// assert_eq!(list_arr, expected);
    519
    0
    pub fn arrays_into_list_array(
    520
    0
        arr: impl IntoIterator<Item = ArrayRef>,
    521
    0
    ) -> Result<ListArray> {
    522
    0
        let arr = arr.into_iter().collect::<Vec<_>>();
    523
    0
        assert_or_internal_err!(!arr.is_empty(), "Cannot wrap empty array into list array");
    524
    525
    0
        let lens = arr.iter().map(|x| x.len()).collect::<Vec<_>>();
    526
        // Assume data type is consistent
    527
    0
        let data_type = arr[0].data_type().to_owned();
    528
    0
        let values = arr.iter().map(|x| x.as_ref()).collect::<Vec<_>>();
    529
    0
        Ok(ListArray::new(
    530
    0
            Arc::new(Field::new_list_field(data_type, true)),
    531
    0
            OffsetBuffer::from_lengths(lens),
    532
    0
            arrow::compute::concat(values.as_slice())?,
    533
    0
            None,
    534
        ))
    535
    0
    }
    536
    537
    /// Helper function to convert a ListArray into a vector of ArrayRefs.
    538
    0
    pub fn list_to_arrays<O: OffsetSizeTrait>(a: &ArrayRef) -> Vec<ArrayRef> {
    539
    0
        a.as_list::<O>().iter().flatten().collect::<Vec<_>>()
    540
    0
    }
    541
    542
    /// Helper function to convert a FixedSizeListArray into a vector of ArrayRefs.
    543
    0
    pub fn fixed_size_list_to_arrays(a: &ArrayRef) -> Vec<ArrayRef> {
    544
    0
        a.as_fixed_size_list().iter().flatten().collect::<Vec<_>>()
    545
    0
    }
    546
    547
    /// Get the base type of a data type.
    548
    ///
    549
    /// Example
    550
    /// ```
    551
    /// use arrow::datatypes::{DataType, Field};
    552
    /// use datafusion_common::utils::base_type;
    553
    /// use std::sync::Arc;
    554
    ///
    555
    /// let data_type =
    556
    ///     DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true)));
    557
    /// assert_eq!(base_type(&data_type), DataType::Int32);
    558
    ///
    559
    /// let data_type = DataType::Int32;
    560
    /// assert_eq!(base_type(&data_type), DataType::Int32);
    561
    /// ```
    562
    0
    pub fn base_type(data_type: &DataType) -> DataType {
    563
    0
        match data_type {
    564
    0
            DataType::List(field)
    565
    0
            | DataType::LargeList(field)
    566
    0
            | DataType::FixedSizeList(field, _) => base_type(field.data_type()),
    567
    0
            _ => data_type.to_owned(),
    568
        }
    569
    0
    }
    570
    571
    /// Information about how to coerce lists.
    572
    #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
    573
    pub enum ListCoercion {
    574
        /// [`DataType::FixedSizeList`] should be coerced to [`DataType::List`].
    575
        FixedSizedListToList,
    576
    }
    577
    578
    /// A helper function to coerce base type in List.
    579
    ///
    580
    /// Example
    581
    /// ```
    582
    /// use arrow::datatypes::{DataType, Field};
    583
    /// use datafusion_common::utils::coerced_type_with_base_type_only;
    584
    /// use std::sync::Arc;
    585
    ///
    586
    /// let data_type = DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true)));
    587
    /// let base_type = DataType::Float64;
    588
    /// let coerced_type = coerced_type_with_base_type_only(&data_type, &base_type, None);
    589
    /// assert_eq!(coerced_type, DataType::List(Arc::new(Field::new_list_field(DataType::Float64, true))));
    590
    0
    pub fn coerced_type_with_base_type_only(
    591
    0
        data_type: &DataType,
    592
    0
        base_type: &DataType,
    593
    0
        array_coercion: Option<&ListCoercion>,
    594
    0
    ) -> DataType {
    595
    0
        match (data_type, array_coercion) {
    596
    0
            (DataType::List(field), _)
    597
    0
            | (DataType::FixedSizeList(field, _), Some(ListCoercion::FixedSizedListToList)) =>
    598
            {
    599
    0
                let field_type = coerced_type_with_base_type_only(
    600
    0
                    field.data_type(),
    601
    0
                    base_type,
    602
    0
                    array_coercion,
    603
                );
    604
    605
    0
                DataType::List(Arc::new(Field::new(
    606
    0
                    field.name(),
    607
    0
                    field_type,
    608
    0
                    field.is_nullable(),
    609
    0
                )))
    610
            }
    611
    0
            (DataType::FixedSizeList(field, len), _) => {
    612
    0
                let field_type = coerced_type_with_base_type_only(
    613
    0
                    field.data_type(),
    614
    0
                    base_type,
    615
    0
                    array_coercion,
    616
                );
    617
    618
    0
                DataType::FixedSizeList(
    619
    0
                    Arc::new(Field::new(field.name(), field_type, field.is_nullable())),
    620
    0
                    *len,
    621
    0
                )
    622
            }
    623
    0
            (DataType::LargeList(field), _) => {
    624
    0
                let field_type = coerced_type_with_base_type_only(
    625
    0
                    field.data_type(),
    626
    0
                    base_type,
    627
    0
                    array_coercion,
    628
                );
    629
    630
    0
                DataType::LargeList(Arc::new(Field::new(
    631
    0
                    field.name(),
    632
    0
                    field_type,
    633
    0
                    field.is_nullable(),
    634
    0
                )))
    635
            }
    636
    637
    0
            _ => base_type.clone(),
    638
        }
    639
    0
    }
    640
    641
    /// Recursively coerce and `FixedSizeList` elements to `List`
    642
    0
    pub fn coerced_fixed_size_list_to_list(data_type: &DataType) -> DataType {
    643
    0
        match data_type {
    644
    0
            DataType::List(field) | DataType::FixedSizeList(field, _) => {
    645
    0
                let field_type = coerced_fixed_size_list_to_list(field.data_type());
    646
    647
    0
                DataType::List(Arc::new(Field::new(
    648
    0
                    field.name(),
    649
    0
                    field_type,
    650
    0
                    field.is_nullable(),
    651
    0
                )))
    652
            }
    653
    0
            DataType::LargeList(field) => {
    654
    0
                let field_type = coerced_fixed_size_list_to_list(field.data_type());
    655
    656
    0
                DataType::LargeList(Arc::new(Field::new(
    657
    0
                    field.name(),
    658
    0
                    field_type,
    659
    0
                    field.is_nullable(),
    660
    0
                )))
    661
            }
    662
    663
    0
            _ => data_type.clone(),
    664
        }
    665
    0
    }
    666
    667
    /// Compute the number of dimensions in a list data type.
    668
    0
    pub fn list_ndims(data_type: &DataType) -> u64 {
    669
    0
        match data_type {
    670
    0
            DataType::List(field)
    671
    0
            | DataType::LargeList(field)
    672
    0
            | DataType::FixedSizeList(field, _) => 1 + list_ndims(field.data_type()),
    673
    0
            _ => 0,
    674
        }
    675
    0
    }
    676
    677
    /// Adopted from strsim-rs for string similarity metrics
    678
    pub mod datafusion_strsim {
    679
        // Source: https://github.com/dguo/strsim-rs/blob/master/src/lib.rs
    680
        // License: https://github.com/dguo/strsim-rs/blob/master/LICENSE
    681
        use std::cmp::min;
    682
        use std::str::Chars;
    683
    684
        struct StringWrapper<'a>(&'a str);
    685
    686
        impl<'b> IntoIterator for &StringWrapper<'b> {
    687
            type Item = char;
    688
            type IntoIter = Chars<'b>;
    689
    690
    0
            fn into_iter(self) -> Self::IntoIter {
    691
    0
                self.0.chars()
    692
    0
            }
    693
        }
    694
    695
        /// Calculates the minimum number of insertions, deletions, and substitutions
    696
        /// required to change one sequence into the other.
    697
    0
        fn generic_levenshtein<'a, 'b, Iter1, Iter2, Elem1, Elem2>(
    698
    0
            a: &'a Iter1,
    699
    0
            b: &'b Iter2,
    700
    0
        ) -> usize
    701
    0
        where
    702
    0
            &'a Iter1: IntoIterator<Item = Elem1>,
    703
    0
            &'b Iter2: IntoIterator<Item = Elem2>,
    704
    0
            Elem1: PartialEq<Elem2>,
    705
        {
    706
    0
            let b_len = b.into_iter().count();
    707
    708
    0
            if a.into_iter().next().is_none() {
    709
    0
                return b_len;
    710
    0
            }
    711
    712
    0
            let mut cache: Vec<usize> = (1..b_len + 1).collect();
    713
    714
    0
            let mut result = 0;
    715
    716
    0
            for (i, a_elem) in a.into_iter().enumerate() {
    717
    0
                result = i + 1;
    718
    0
                let mut distance_b = i;
    719
    720
    0
                for (j, b_elem) in b.into_iter().enumerate() {
    721
    0
                    let cost = if a_elem == b_elem { 0usize } else { 1usize };
    722
    0
                    let distance_a = distance_b + cost;
    723
    0
                    distance_b = cache[j];
    724
    0
                    result = min(result + 1, min(distance_a, distance_b + 1));
    725
    0
                    cache[j] = result;
    726
                }
    727
            }
    728
    729
    0
            result
    730
    0
        }
    731
    732
        /// Calculates the minimum number of insertions, deletions, and substitutions
    733
        /// required to change one string into the other.
    734
        ///
    735
        /// ```
    736
        /// use datafusion_common::utils::datafusion_strsim::levenshtein;
    737
        ///
    738
        /// assert_eq!(3, levenshtein("kitten", "sitting"));
    739
        /// ```
    740
    0
        pub fn levenshtein(a: &str, b: &str) -> usize {
    741
    0
            generic_levenshtein(&StringWrapper(a), &StringWrapper(b))
    742
    0
        }
    743
    744
        /// Calculates the normalized Levenshtein distance between two strings.
    745
        /// The normalized distance is a value between 0.0 and 1.0, where 1.0 indicates
    746
        /// that the strings are identical and 0.0 indicates no similarity.
    747
        ///
    748
        /// ```
    749
        /// use datafusion_common::utils::datafusion_strsim::normalized_levenshtein;
    750
        ///
    751
        /// assert!((normalized_levenshtein("kitten", "sitting") - 0.57142).abs() < 0.00001);
    752
        ///
    753
        /// assert!(normalized_levenshtein("", "second").abs() < 0.00001);
    754
        ///
    755
        /// assert!((normalized_levenshtein("kitten", "sitten") - 0.833).abs() < 0.001);
    756
        /// ```
    757
    0
        pub fn normalized_levenshtein(a: &str, b: &str) -> f64 {
    758
    0
            if a.is_empty() && b.is_empty() {
    759
    0
                return 1.0;
    760
    0
            }
    761
    0
            1.0 - (levenshtein(a, b) as f64)
    762
    0
                / (a.chars().count().max(b.chars().count()) as f64)
    763
    0
        }
    764
    }
    765
    766
    /// Merges collections `first` and `second`, removes duplicates and sorts the
    767
    /// result, returning it as a [`Vec`].
    768
    0
    pub fn merge_and_order_indices<T: Borrow<usize>, S: Borrow<usize>>(
    769
    0
        first: impl IntoIterator<Item = T>,
    770
    0
        second: impl IntoIterator<Item = S>,
    771
    0
    ) -> Vec<usize> {
    772
    0
        let mut result: Vec<_> = first
    773
    0
            .into_iter()
    774
    0
            .map(|e| *e.borrow())
    775
    0
            .chain(second.into_iter().map(|e| *e.borrow()))
    776
    0
            .collect::<HashSet<_>>()
    777
    0
            .into_iter()
    778
    0
            .collect();
    779
    0
        result.sort();
    780
    0
        result
    781
    0
    }
    782
    783
    /// Calculates the set difference between sequences `first` and `second`,
    784
    /// returning the result as a [`Vec`]. Preserves the ordering of `first`.
    785
    0
    pub fn set_difference<T: Borrow<usize>, S: Borrow<usize>>(
    786
    0
        first: impl IntoIterator<Item = T>,
    787
    0
        second: impl IntoIterator<Item = S>,
    788
    0
    ) -> Vec<usize> {
    789
    0
        let set: HashSet<_> = second.into_iter().map(|e| *e.borrow()).collect();
    790
    0
        first
    791
    0
            .into_iter()
    792
    0
            .map(|e| *e.borrow())
    793
    0
            .filter(|e| !set.contains(e))
    794
    0
            .collect()
    795
    0
    }
    796
    797
    /// Find indices of each element in `targets` inside `items`. If one of the
    798
    /// elements is absent in `items`, returns an error.
    799
    0
    pub fn find_indices<T: PartialEq, S: Borrow<T>>(
    800
    0
        items: &[T],
    801
    0
        targets: impl IntoIterator<Item = S>,
    802
    0
    ) -> Result<Vec<usize>> {
    803
    0
        targets
    804
    0
            .into_iter()
    805
    0
            .map(|target| items.iter().position(|e| target.borrow().eq(e)))
    806
    0
            .collect::<Option<_>>()
    807
    0
            .ok_or_else(|| _exec_datafusion_err!("Target not found"))
    808
    0
    }
    809
    810
    /// Transposes the given vector of vectors.
    811
    0
    pub fn transpose<T>(original: Vec<Vec<T>>) -> Vec<Vec<T>> {
    812
    0
        match original.as_slice() {
    813
    0
            [] => vec![],
    814
    0
            [first, ..] => {
    815
    0
                let mut result = (0..first.len()).map(|_| vec![]).collect::<Vec<_>>();
    816
    0
                for row in original {
    817
    0
                    for (item, transposed_row) in row.into_iter().zip(&mut result) {
    818
    0
                        transposed_row.push(item);
    819
    0
                    }
    820
                }
    821
    0
                result
    822
            }
    823
        }
    824
    0
    }
    825
    826
    /// Computes the `skip` and `fetch` parameters of a single limit that would be
    827
    /// equivalent to two consecutive limits with the given `skip`/`fetch` parameters.
    828
    ///
    829
    /// There are multiple cases to consider:
    830
    ///
    831
    /// # Case 0: Parent and child are disjoint (`child_fetch <= skip`).
    832
    ///
    833
    /// ```text
    834
    ///   Before merging:
    835
    ///                     |........skip........|---fetch-->|     Parent limit
    836
    ///    |...child_skip...|---child_fetch-->|                    Child limit
    837
    /// ```
    838
    ///
    839
    ///   After merging:
    840
    /// ```text
    841
    ///    |.........(child_skip + skip).........|
    842
    /// ```
    843
    ///
    844
    /// # Case 1: Parent is beyond child's range (`skip < child_fetch <= skip + fetch`).
    845
    ///
    846
    ///   Before merging:
    847
    /// ```text
    848
    ///                     |...skip...|------------fetch------------>|   Parent limit
    849
    ///    |...child_skip...|-------------child_fetch------------>|       Child limit
    850
    /// ```
    851
    ///
    852
    ///   After merging:
    853
    /// ```text
    854
    ///    |....(child_skip + skip)....|---(child_fetch - skip)-->|
    855
    /// ```
    856
    ///
    857
    ///  # Case 2: Parent is within child's range (`skip + fetch < child_fetch`).
    858
    ///
    859
    ///   Before merging:
    860
    /// ```text
    861
    ///                     |...skip...|---fetch-->|                   Parent limit
    862
    ///    |...child_skip...|-------------child_fetch------------>|    Child limit
    863
    /// ```
    864
    ///
    865
    ///   After merging:
    866
    /// ```text
    867
    ///    |....(child_skip + skip)....|---fetch-->|
    868
    /// ```
    869
    0
    pub fn combine_limit(
    870
    0
        parent_skip: usize,
    871
    0
        parent_fetch: Option<usize>,
    872
    0
        child_skip: usize,
    873
    0
        child_fetch: Option<usize>,
    874
    0
    ) -> (usize, Option<usize>) {
    875
    0
        let combined_skip = child_skip.saturating_add(parent_skip);
    876
    877
    0
        let combined_fetch = match (parent_fetch, child_fetch) {
    878
    0
            (Some(parent_fetch), Some(child_fetch)) => {
    879
    0
                Some(min(parent_fetch, child_fetch.saturating_sub(parent_skip)))
    880
            }
    881
    0
            (Some(parent_fetch), None) => Some(parent_fetch),
    882
    0
            (None, Some(child_fetch)) => Some(child_fetch.saturating_sub(parent_skip)),
    883
    0
            (None, None) => None,
    884
        };
    885
    886
    0
        (combined_skip, combined_fetch)
    887
    0
    }
    888
    889
    /// Returns the estimated number of threads available for parallel execution.
    890
    ///
    891
    /// This is a wrapper around `std::thread::available_parallelism`, providing a default value
    892
    /// of `1` if the system's parallelism cannot be determined.
    893
    0
    pub fn get_available_parallelism() -> usize {
    894
    0
        available_parallelism()
    895
    0
            .unwrap_or(NonZero::new(1).expect("literal value `1` shouldn't be zero"))
    896
    0
            .get()
    897
    0
    }
    898
    899
    /// Converts a collection of function arguments into a fixed-size array of length N
    900
    /// producing a reasonable error message in case of unexpected number of arguments.
    901
    ///
    902
    /// # Example
    903
    /// ```
    904
    /// # use datafusion_common::Result;
    905
    /// # use datafusion_common::utils::take_function_args;
    906
    /// # use datafusion_common::ScalarValue;
    907
    /// fn my_function(args: &[ScalarValue]) -> Result<()> {
    908
    ///     // function expects 2 args, so create a 2-element array
    909
    ///     let [arg1, arg2] = take_function_args("my_function", args)?;
    910
    ///     // ... do stuff..
    911
    ///     Ok(())
    912
    /// }
    913
    ///
    914
    /// // Calling the function with 1 argument produces an error:
    915
    /// let args = vec![ScalarValue::Int32(Some(10))];
    916
    /// let err = my_function(&args).unwrap_err();
    917
    /// assert_eq!(
    918
    ///     err.to_string(),
    919
    ///     "Execution error: my_function function requires 2 arguments, got 1"
    920
    /// );
    921
    /// // Calling the function with 2 arguments works great
    922
    /// let args = vec![ScalarValue::Int32(Some(10)), ScalarValue::Int32(Some(20))];
    923
    /// my_function(&args).unwrap();
    924
    /// ```
    925
    0
    pub fn take_function_args<const N: usize, T>(
    926
    0
        function_name: &str,
    927
    0
        args: impl IntoIterator<Item = T>,
    928
    0
    ) -> Result<[T; N]> {
    929
    0
        let args = args.into_iter().collect::<Vec<_>>();
    930
    0
        args.try_into().map_err(|v: Vec<T>| {
    931
    0
            _exec_datafusion_err!(
    932
    0
                "{} function requires {} {}, got {}",
    933
                function_name,
    934
                N,
    935
    0
                if N == 1 { "argument" } else { "arguments" },
    936
    0
                v.len()
    937
            )
    938
    0
        })
    939
    0
    }
    940
    941
    #[cfg(test)]
    942
    mod tests {
    943
        use super::*;
    944
        use crate::ScalarValue::Null;
    945
        use arrow::array::Float64Array;
    946
        use sqlparser::ast::Ident;
    947
        use sqlparser::tokenizer::Span;
    948
    949
        #[test]
    950
        fn test_bisect_linear_left_and_right() -> Result<()> {
    951
            let arrays: Vec<ArrayRef> = vec![
    952
                Arc::new(Float64Array::from(vec![5.0, 7.0, 8.0, 9., 10.])),
    953
                Arc::new(Float64Array::from(vec![2.0, 3.0, 3.0, 4.0, 5.0])),
    954
                Arc::new(Float64Array::from(vec![5.0, 7.0, 8.0, 10., 11.0])),
    955
                Arc::new(Float64Array::from(vec![15.0, 13.0, 8.0, 5., 0.0])),
    956
            ];
    957
            let search_tuple: Vec<ScalarValue> = vec![
    958
                ScalarValue::Float64(Some(8.0)),
    959
                ScalarValue::Float64(Some(3.0)),
    960
                ScalarValue::Float64(Some(8.0)),
    961
                ScalarValue::Float64(Some(8.0)),
    962
            ];
    963
            let ords = [
    964
                SortOptions {
    965
                    descending: false,
    966
                    nulls_first: true,
    967
                },
    968
                SortOptions {
    969
                    descending: false,
    970
                    nulls_first: true,
    971
                },
    972
                SortOptions {
    973
                    descending: false,
    974
                    nulls_first: true,
    975
                },
    976
                SortOptions {
    977
                    descending: true,
    978
                    nulls_first: true,
    979
                },
    980
            ];
    981
            let res = bisect::<true>(&arrays, &search_tuple, &ords)?;
    982
            assert_eq!(res, 2);
    983
            let res = bisect::<false>(&arrays, &search_tuple, &ords)?;
    984
            assert_eq!(res, 3);
    985
            let res = linear_search::<true>(&arrays, &search_tuple, &ords)?;
    986
            assert_eq!(res, 2);
    987
            let res = linear_search::<false>(&arrays, &search_tuple, &ords)?;
    988
            assert_eq!(res, 3);
    989
            Ok(())
    990
        }
    991
    992
        #[test]
    993
        fn vector_ord() {
    994
            assert!(vec![1, 0, 0, 0, 0, 0, 0, 1] < vec![1, 0, 0, 0, 0, 0, 0, 2]);
    995
            assert!(vec![1, 0, 0, 0, 0, 0, 1, 1] > vec![1, 0, 0, 0, 0, 0, 0, 2]);
    996
            assert!(
    997
                vec![
    998
                    ScalarValue::Int32(Some(2)),
    999
                    Null,
    1000
                    ScalarValue::Int32(Some(0)),
    1001
                ] < vec![
    1002
                    ScalarValue::Int32(Some(2)),
    1003
                    Null,
    1004
                    ScalarValue::Int32(Some(1)),
    1005
                ]
    1006
            );
    1007
            assert!(
    1008
                vec![
    1009
                    ScalarValue::Int32(Some(2)),
    1010
                    ScalarValue::Int32(None),
    1011
                    ScalarValue::Int32(Some(0)),
    1012
                ] < vec![
    1013
                    ScalarValue::Int32(Some(2)),
    1014
                    ScalarValue::Int32(None),
    1015
                    ScalarValue::Int32(Some(1)),
    1016
                ]
    1017
            );
    1018
        }
    1019
    1020
        #[test]
    1021
        fn ord_same_type() {
    1022
            assert!((ScalarValue::Int32(Some(2)) < ScalarValue::Int32(Some(3))));
    1023
        }
    1024
    1025
        #[test]
    1026
        fn test_bisect_linear_left_and_right_diff_sort() -> Result<()> {
    1027
            // Descending, left
    1028
            let arrays: Vec<ArrayRef> =
    1029
                vec![Arc::new(Float64Array::from(vec![4.0, 3.0, 2.0, 1.0, 0.0]))];
    1030
            let search_tuple: Vec<ScalarValue> = vec![ScalarValue::Float64(Some(4.0))];
    1031
            let ords = [SortOptions {
    1032
                descending: true,
    1033
                nulls_first: true,
    1034
            }];
    1035
            let res = bisect::<true>(&arrays, &search_tuple, &ords)?;
    1036
            assert_eq!(res, 0);
    1037
            let res = linear_search::<true>(&arrays, &search_tuple, &ords)?;
    1038
            assert_eq!(res, 0);
    1039
    1040
            // Descending, right
    1041
            let arrays: Vec<ArrayRef> =
    1042
                vec![Arc::new(Float64Array::from(vec![4.0, 3.0, 2.0, 1.0, 0.0]))];
    1043
            let search_tuple: Vec<ScalarValue> = vec![ScalarValue::Float64(Some(4.0))];
    1044
            let ords = [SortOptions {
    1045
                descending: true,
    1046
                nulls_first: true,
    1047
            }];
    1048
            let res = bisect::<false>(&arrays, &search_tuple, &ords)?;
    1049
            assert_eq!(res, 1);
    1050
            let res = linear_search::<false>(&arrays, &search_tuple, &ords)?;
    1051
            assert_eq!(res, 1);
    1052
    1053
            // Ascending, left
    1054
            let arrays: Vec<ArrayRef> =
    1055
                vec![Arc::new(Float64Array::from(vec![5.0, 7.0, 8.0, 9., 10.]))];
    1056
            let search_tuple: Vec<ScalarValue> = vec![ScalarValue::Float64(Some(7.0))];
    1057
            let ords = [SortOptions {
    1058
                descending: false,
    1059
                nulls_first: true,
    1060
            }];
    1061
            let res = bisect::<true>(&arrays, &search_tuple, &ords)?;
    1062
            assert_eq!(res, 1);
    1063
            let res = linear_search::<true>(&arrays, &search_tuple, &ords)?;
    1064
            assert_eq!(res, 1);
    1065
    1066
            // Ascending, right
    1067
            let arrays: Vec<ArrayRef> =
    1068
                vec![Arc::new(Float64Array::from(vec![5.0, 7.0, 8.0, 9., 10.]))];
    1069
            let search_tuple: Vec<ScalarValue> = vec![ScalarValue::Float64(Some(7.0))];
    1070
            let ords = [SortOptions {
    1071
                descending: false,
    1072
                nulls_first: true,
    1073
            }];
    1074
            let res = bisect::<false>(&arrays, &search_tuple, &ords)?;
    1075
            assert_eq!(res, 2);
    1076
            let res = linear_search::<false>(&arrays, &search_tuple, &ords)?;
    1077
            assert_eq!(res, 2);
    1078
    1079
            let arrays: Vec<ArrayRef> = vec![
    1080
                Arc::new(Float64Array::from(vec![5.0, 7.0, 8.0, 8.0, 9., 10.])),
    1081
                Arc::new(Float64Array::from(vec![10.0, 9.0, 8.0, 7.5, 7., 6.])),
    1082
            ];
    1083
            let search_tuple: Vec<ScalarValue> = vec![
    1084
                ScalarValue::Float64(Some(8.0)),
    1085
                ScalarValue::Float64(Some(8.0)),
    1086
            ];
    1087
            let ords = [
    1088
                SortOptions {
    1089
                    descending: false,
    1090
                    nulls_first: true,
    1091
                },
    1092
                SortOptions {
    1093
                    descending: true,
    1094
                    nulls_first: true,
    1095
                },
    1096
            ];
    1097
            let res = bisect::<false>(&arrays, &search_tuple, &ords)?;
    1098
            assert_eq!(res, 3);
    1099
            let res = linear_search::<false>(&arrays, &search_tuple, &ords)?;
    1100
            assert_eq!(res, 3);
    1101
    1102
            let res = bisect::<true>(&arrays, &search_tuple, &ords)?;
    1103
            assert_eq!(res, 2);
    1104
            let res = linear_search::<true>(&arrays, &search_tuple, &ords)?;
    1105
            assert_eq!(res, 2);
    1106
            Ok(())
    1107
        }
    1108
    1109
        #[test]
    1110
        fn test_evaluate_partition_ranges() -> Result<()> {
    1111
            let arrays: Vec<ArrayRef> = vec![
    1112
                Arc::new(Float64Array::from(vec![1.0, 1.0, 1.0, 2.0, 2.0, 2.0])),
    1113
                Arc::new(Float64Array::from(vec![4.0, 4.0, 3.0, 2.0, 1.0, 1.0])),
    1114
            ];
    1115
            let n_row = arrays[0].len();
    1116
            let options: Vec<SortOptions> = vec![
    1117
                SortOptions {
    1118
                    descending: false,
    1119
                    nulls_first: false,
    1120
                },
    1121
                SortOptions {
    1122
                    descending: true,
    1123
                    nulls_first: false,
    1124
                },
    1125
            ];
    1126
            let sort_columns = arrays
    1127
                .into_iter()
    1128
                .zip(options)
    1129
                .map(|(values, options)| SortColumn {
    1130
                    values,
    1131
                    options: Some(options),
    1132
                })
    1133
                .collect::<Vec<_>>();
    1134
            let ranges = evaluate_partition_ranges(n_row, &sort_columns)?;
    1135
            assert_eq!(ranges.len(), 4);
    1136
            assert_eq!(ranges[0], Range { start: 0, end: 2 });
    1137
            assert_eq!(ranges[1], Range { start: 2, end: 3 });
    1138
            assert_eq!(ranges[2], Range { start: 3, end: 4 });
    1139
            assert_eq!(ranges[3], Range { start: 4, end: 6 });
    1140
            Ok(())
    1141
        }
    1142
    1143
        #[cfg(feature = "sql")]
    1144
        #[test]
    1145
        fn test_quote_identifier() -> Result<()> {
    1146
            let cases = vec![
    1147
                ("foo", r#"foo"#),
    1148
                ("_foo", r#"_foo"#),
    1149
                ("foo_bar", r#"foo_bar"#),
    1150
                ("foo-bar", r#""foo-bar""#),
    1151
                // name itself has a period, needs to be quoted
    1152
                ("foo.bar", r#""foo.bar""#),
    1153
                ("Foo", r#""Foo""#),
    1154
                ("Foo.Bar", r#""Foo.Bar""#),
    1155
                // name starting with a number needs to be quoted
    1156
                ("test1", r#"test1"#),
    1157
                ("1test", r#""1test""#),
    1158
            ];
    1159
    1160
            for (identifier, quoted_identifier) in cases {
    1161
                println!("input: \n{identifier}\nquoted_identifier:\n{quoted_identifier}");
    1162
    1163
                assert_eq!(quote_identifier(identifier), quoted_identifier);
    1164
    1165
                // When parsing the quoted identifier, it should be a
    1166
                // a single identifier without normalization, and not in multiple parts
    1167
                let quote_style = if quoted_identifier.starts_with('"') {
    1168
                    Some('"')
    1169
                } else {
    1170
                    None
    1171
                };
    1172
    1173
                let expected_parsed = vec![Ident {
    1174
                    value: identifier.to_string(),
    1175
                    quote_style,
    1176
                    span: Span::empty(),
    1177
                }];
    1178
    1179
                assert_eq!(
    1180
                    parse_identifiers(quoted_identifier).unwrap(),
    1181
                    expected_parsed
    1182
                );
    1183
            }
    1184
    1185
            Ok(())
    1186
        }
    1187
    1188
        #[test]
    1189
        fn test_get_at_indices() -> Result<()> {
    1190
            let in_vec = vec![1, 2, 3, 4, 5, 6, 7];
    1191
            assert_eq!(get_at_indices(&in_vec, [0, 2])?, vec![1, 3]);
    1192
            assert_eq!(get_at_indices(&in_vec, [4, 2])?, vec![5, 3]);
    1193
            // 7 is outside the range
    1194
            assert!(get_at_indices(&in_vec, [7]).is_err());
    1195
            Ok(())
    1196
        }
    1197
    1198
        #[test]
    1199
        fn test_longest_consecutive_prefix() {
    1200
            assert_eq!(longest_consecutive_prefix([0, 3, 4]), 1);
    1201
            assert_eq!(longest_consecutive_prefix([0, 1, 3, 4]), 2);
    1202
            assert_eq!(longest_consecutive_prefix([0, 1, 2, 3, 4]), 5);
    1203
            assert_eq!(longest_consecutive_prefix([1, 2, 3, 4]), 0);
    1204
        }
    1205
    1206
        #[test]
    1207
        fn test_merge_and_order_indices() {
    1208
            assert_eq!(
    1209
                merge_and_order_indices([0, 3, 4], [1, 3, 5]),
    1210
                vec![0, 1, 3, 4, 5]
    1211
            );
    1212
            // Result should be ordered, even if inputs are not
    1213
            assert_eq!(
    1214
                merge_and_order_indices([3, 0, 4], [5, 1, 3]),
    1215
                vec![0, 1, 3, 4, 5]
    1216
            );
    1217
        }
    1218
    1219
        #[test]
    1220
        fn test_set_difference() {
    1221
            assert_eq!(set_difference([0, 3, 4], [1, 2]), vec![0, 3, 4]);
    1222
            assert_eq!(set_difference([0, 3, 4], [1, 2, 4]), vec![0, 3]);
    1223
            // return value should have same ordering with the in1
    1224
            assert_eq!(set_difference([3, 4, 0], [1, 2, 4]), vec![3, 0]);
    1225
            assert_eq!(set_difference([0, 3, 4], [4, 1, 2]), vec![0, 3]);
    1226
            assert_eq!(set_difference([3, 4, 0], [4, 1, 2]), vec![3, 0]);
    1227
        }
    1228
    1229
        #[test]
    1230
        fn test_find_indices() -> Result<()> {
    1231
            assert_eq!(find_indices(&[0, 3, 4], [0, 3, 4])?, vec![0, 1, 2]);
    1232
            assert_eq!(find_indices(&[0, 3, 4], [0, 4, 3])?, vec![0, 2, 1]);
    1233
            assert_eq!(find_indices(&[3, 0, 4], [0, 3])?, vec![1, 0]);
    1234
            assert!(find_indices(&[0, 3], [0, 3, 4]).is_err());
    1235
            assert!(find_indices(&[0, 3, 4], [0, 2]).is_err());
    1236
            Ok(())
    1237
        }
    1238
    1239
        #[test]
    1240
        fn test_transpose() -> Result<()> {
    1241
            let in_data = vec![vec![1, 2, 3], vec![4, 5, 6]];
    1242
            let transposed = transpose(in_data);
    1243
            let expected = vec![vec![1, 4], vec![2, 5], vec![3, 6]];
    1244
            assert_eq!(expected, transposed);
    1245
            Ok(())
    1246
        }
    1247
    }