// SPDX-License-Identifier: Apache-2.0

// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership

// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements

//! This module provides the bisect function, which implements binary search.

pub(crate) mod aggregate;

pub mod expr;

pub mod hex;

pub mod memory;

pub mod proxy;

pub mod string_utils;

use crate::assert_or_internal_err;

use crate::error::{

_exec_datafusion_err, _exec_err, _internal_datafusion_err, _internal_err,

_plan_datafusion_err,

};

use crate::hash_utils::RandomState;

use crate::{Result, ScalarValue};

use arrow::array::{

Array, ArrayRef, FixedSizeListArray, LargeListArray, ListArray, OffsetSizeTrait,

cast::AsArray,

};

use arrow::array::{

ArrowPrimitiveType, BooleanArray, Datum, DictionaryArray, GenericListArray,

Int32Array, Int64Array, MutableArrayData, PrimitiveArray, layout, make_array,

};

use arrow::array::{LargeListViewArray, ListViewArray};

use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer};

use arrow::compute::kernels::cmp::eq;

use arrow::compute::kernels::length::length;

use arrow::compute::{SortColumn, SortOptions, nullif, partition, take};

use arrow::datatypes::{

ArrowDictionaryKeyType, ArrowNativeType, DataType, Field, Int32Type, Int64Type,

Metadata, SchemaRef,

};

use arrow::downcast_dictionary_array;

use arrow_select::dictionary::garbage_collect_any_dictionary;

#[cfg(feature = "sql")]

use sqlparser::{ast::Ident, dialect::GenericDialect, parser::Parser};

use std::borrow::{Borrow, Cow};

use std::cmp::{Ordering, min};

use std::collections::{HashMap, HashSet};

use std::iter::repeat_n;

use std::num::NonZero;

use std::ops::Range;

use std::sync::{Arc, LazyLock};

use std::thread::available_parallelism;

/// Applies an optional projection to a [`SchemaRef`], returning the

/// projected schema

///

/// Example:

/// ```

/// use arrow::datatypes::{DataType, Field, Schema, SchemaRef};

/// use datafusion_common::project_schema;

///

/// // Schema with columns 'a', 'b', and 'c'

/// let schema = SchemaRef::new(Schema::new(vec![

/// Field::new("a", DataType::Int32, true),

/// Field::new("b", DataType::Int64, true),

/// Field::new("c", DataType::Utf8, true),

/// ]));

///

/// // Pick columns 'c' and 'b'

/// let projection = Some(vec![2, 1]);

/// let projected_schema = project_schema(&schema, projection.as_ref()).unwrap();

///

/// let expected_schema = SchemaRef::new(Schema::new(vec![

/// Field::new("c", DataType::Utf8, true),

/// Field::new("b", DataType::Int64, true),

/// ]));

///

/// assert_eq!(projected_schema, expected_schema);

/// ```

pub fn project_schema<T: AsRef<[usize]> + ?Sized>(

schema: &SchemaRef,

projection: Option<&T>,

) -> Result<SchemaRef> {

let schema = match projection {

Some(columns) => Arc::new(schema.project(columns.as_ref())?),

None => Arc::clone(schema),

};

Ok(schema)

}

/// Extracts a row at the specified index from a set of columns and stores it in the provided buffer.

pub fn extract_row_at_idx_to_buf(

columns: &[ArrayRef],

idx: usize,

buf: &mut Vec<ScalarValue>,

) -> Result<()> {

buf.clear();

let iter = columns

.iter()

.map(|arr| ScalarValue::try_from_array(arr, idx));

for v in iter.into_iter() {

buf.push(v?);

}

Ok(())

}

/// Given column vectors, returns row at `idx`.

pub fn get_row_at_idx(columns: &[ArrayRef], idx: usize) -> Result<Vec<ScalarValue>> {

columns

.iter()

.map(|arr| ScalarValue::try_from_array(arr, idx))

.collect()

}

/// This function compares two tuples depending on the given sort options.

pub fn compare_rows(

x: &[ScalarValue],

y: &[ScalarValue],

sort_options: &[SortOptions],

) -> Result<Ordering> {

let zip_it = x.iter().zip(y.iter()).zip(sort_options.iter());

// Preserving lexical ordering.

for ((lhs, rhs), sort_options) in zip_it {

// Consider all combinations of NULLS FIRST/LAST and ASC/DESC configurations.

let result = match (lhs.is_null(), rhs.is_null(), sort_options.nulls_first) {

(true, false, false) | (false, true, true) => Ordering::Greater,

(true, false, true) | (false, true, false) => Ordering::Less,

(false, false, _) => {

if sort_options.descending {

rhs.try_cmp(lhs)?

} else {

lhs.try_cmp(rhs)?

}

}

(true, true, _) => continue,

};

if result != Ordering::Equal {

return Ok(result);

}

}

Ok(Ordering::Equal)

}

/// This function searches for a tuple of given values (`target`) among the given

/// rows (`item_columns`) using the bisection algorithm. It assumes that `item_columns`

/// is sorted according to `sort_options` and returns the insertion index of `target`.

/// Template argument `SIDE` being `true`/`false` means left/right insertion.

pub fn bisect<const SIDE: bool>(

item_columns: &[ArrayRef],

target: &[ScalarValue],

sort_options: &[SortOptions],

) -> Result<usize> {

let low: usize = 0;

let high: usize = item_columns

.first()

.ok_or_else(|| _internal_datafusion_err!("Column array shouldn't be empty"))?

.len();

let compare_fn = |current: &[ScalarValue], target: &[ScalarValue]| {

let cmp = compare_rows(current, target, sort_options)?;

Ok(if SIDE { cmp.is_lt() } else { cmp.is_le() })

};

find_bisect_point(item_columns, target, compare_fn, low, high)

}

/// This function searches for a tuple of given values (`target`) among a slice of

/// the given rows (`item_columns`) using the bisection algorithm. The slice starts

/// at the index `low` and ends at the index `high`. The boolean-valued function

/// `compare_fn` specifies whether we bisect on the left (by returning `false`),

/// or on the right (by returning `true`) when we compare the target value with

/// the current value as we iteratively bisect the input.

pub fn find_bisect_point<F>(

item_columns: &[ArrayRef],

target: &[ScalarValue],

compare_fn: F,

mut low: usize,

mut high: usize,

) -> Result<usize>

where

F: Fn(&[ScalarValue], &[ScalarValue]) -> Result<bool>,

{

let mut val = Vec::with_capacity(item_columns.len());

while low < high {

let mid = ((high - low) / 2) + low;

extract_row_at_idx_to_buf(item_columns, mid, &mut val)?;

if compare_fn(&val, target)? {

low = mid + 1;

} else {

high = mid;

}

}

Ok(low)

}

/// This function searches for a tuple of given values (`target`) among the given

/// rows (`item_columns`) via a linear scan. It assumes that `item_columns` is sorted

/// according to `sort_options` and returns the insertion index of `target`.

/// Template argument `SIDE` being `true`/`false` means left/right insertion.

pub fn linear_search<const SIDE: bool>(

item_columns: &[ArrayRef],

target: &[ScalarValue],

sort_options: &[SortOptions],

) -> Result<usize> {

let low: usize = 0;

let high: usize = item_columns

.first()

.ok_or_else(|| _internal_datafusion_err!("Column array shouldn't be empty"))?

.len();

let compare_fn = |current: &[ScalarValue], target: &[ScalarValue]| {

let cmp = compare_rows(current, target, sort_options)?;

Ok(if SIDE { cmp.is_lt() } else { cmp.is_le() })

};

search_in_slice(item_columns, target, compare_fn, low, high)

}

/// This function searches for a tuple of given values (`target`) among a slice of

/// the given rows (`item_columns`) via a linear scan. The slice starts at the index

/// `low` and ends at the index `high`. The boolean-valued function `compare_fn`

/// specifies the stopping criterion.

pub fn search_in_slice<F>(

item_columns: &[ArrayRef],

target: &[ScalarValue],

compare_fn: F,

mut low: usize,

high: usize,

) -> Result<usize>

where

F: Fn(&[ScalarValue], &[ScalarValue]) -> Result<bool>,

{

let mut val = Vec::new();

while low < high {

extract_row_at_idx_to_buf(item_columns, low, &mut val)?;

if !compare_fn(&val, target)? {

break;

}

low += 1;

}

Ok(low)

}

/// Given a list of 0 or more already sorted columns, finds the

/// partition ranges that would partition equally across columns.

///

/// See [`partition`] for more details.

pub fn evaluate_partition_ranges(

num_rows: usize,

partition_columns: &[SortColumn],

) -> Result<Vec<Range<usize>>> {

Ok(if partition_columns.is_empty() {

vec![Range {

start: 0,

end: num_rows,

}]

} else {

let cols: Vec<_> = partition_columns

.iter()

.map(|x| Arc::clone(&x.values))

.collect();

partition(&cols)?.ranges()

})

}

/// Wraps identifier string in double quotes, escaping any double quotes in

/// the identifier by replacing it with two double quotes

///

/// e.g. identifier `tab.le"name` becomes `"tab.le""name"`

pub fn quote_identifier(s: &str) -> Cow<'_, str> {

if needs_quotes(s) {

Cow::Owned(format!("\"{}\"", s.replace('"', "\"\"")))

} else {

Cow::Borrowed(s)

}

}

/// returns true if this identifier needs quotes

fn needs_quotes(s: &str) -> bool {

let mut chars = s.chars();

// first char can not be a number unless escaped

if let Some(first_char) = chars.next()

&& !(first_char.is_ascii_lowercase() || first_char == '_')

{

return true;

}

!chars.all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_')

}

#[cfg(feature = "sql")]

pub(crate) fn parse_identifiers(s: &str) -> Result<Vec<Ident>> {

let dialect = GenericDialect;

let mut parser = Parser::new(&dialect).try_with_sql(s)?;

let idents = parser.parse_multipart_identifier()?;

Ok(idents)

}

/// Parse a string into a vector of identifiers.

///

/// Note: If ignore_case is false, the string will be normalized to lowercase.

#[cfg(feature = "sql")]

pub(crate) fn parse_identifiers_normalized(s: &str, ignore_case: bool) -> Vec<String> {

parse_identifiers(s)

.unwrap_or_default()

.into_iter()

.map(|id| match id.quote_style {

Some(_) => id.value,

None if ignore_case => id.value,

_ => id.value.to_ascii_lowercase(),

})

.collect::<Vec<_>>()

}

#[cfg(not(feature = "sql"))]

pub(crate) fn parse_identifiers(s: &str) -> Result<Vec<String>> {

let mut result = Vec::new();

let mut current = String::new();

let mut in_quotes = false;

for ch in s.chars() {

match ch {

'"' => {

in_quotes = !in_quotes;

current.push(ch);

}

'.' if !in_quotes => {

result.push(current.clone());

current.clear();

}

_ => {

current.push(ch);

}

}

}

// Push the last part if it's not empty

if !current.is_empty() {

result.push(current);

}

Ok(result)

}

#[cfg(not(feature = "sql"))]

pub(crate) fn parse_identifiers_normalized(s: &str, ignore_case: bool) -> Vec<String> {

parse_identifiers(s)

.unwrap_or_default()

.into_iter()

.map(|id| {

let is_double_quoted = if id.len() > 2 {

let mut chars = id.chars();

chars.next() == Some('"') && chars.last() == Some('"')

} else {

false

};

if is_double_quoted {

id[1..id.len() - 1].replace("\"\"", "\"")

} else if ignore_case {

id

} else {

id.to_ascii_lowercase()

}

})

.collect::<Vec<_>>()

}

/// This function "takes" the elements at `indices` from the slice `items`.

pub fn get_at_indices<T: Clone, I: Borrow<usize>>(

items: &[T],

indices: impl IntoIterator<Item = I>,

) -> Result<Vec<T>> {

indices

.into_iter()

.map(|idx| items.get(*idx.borrow()).cloned())

.collect::<Option<Vec<T>>>()

.ok_or_else(|| {

_exec_datafusion_err!("Expects indices to be in the range of searched vector")

})

}

/// This function finds the longest prefix of the form 0, 1, 2, ... within the

/// collection `sequence`. Examples:

/// - For 0, 1, 2, 4, 5; we would produce 3, meaning 0, 1, 2 is the longest satisfying

/// prefix.

/// - For 1, 2, 3, 4; we would produce 0, meaning there is no such prefix.

pub fn longest_consecutive_prefix<T: Borrow<usize>>(

sequence: impl IntoIterator<Item = T>,

) -> usize {

let mut count = 0;

for item in sequence {

if !count.eq(item.borrow()) {

break;

}

count += 1;

}

count

}

/// Splits `vec` at index `n`, returning the first `n` elements and leaving the

/// remaining `vec.len() - n` elements in `vec`.

///

/// Allocates for whichever side is smaller, so the new allocation is

/// `min(n, vec.len() - n)` rather than always `n` (as `vec.drain(0..n).collect()`

/// would). This matters when the split emits a prefix under memory pressure,

/// where `n` can be close to `vec.len()`.

pub fn split_vec_min_alloc<T>(vec: &mut Vec<T>, n: usize) -> Vec<T> {

if n * 2 <= vec.len() {

vec.drain(0..n).collect()

} else {

let remaining = vec.split_off(n);

std::mem::replace(vec, remaining)

}

}

/// Applies `f` to the values of the dictionary `array` and keeps its keys.

///

/// Slices, filters and repartitioned batches keep the whole values array of

/// the dictionary they came from, so running `f` on every value can cost far

/// more than the batch has rows. `f` never runs on more than

/// [`MAX_DICTIONARY_VALUES_PER_ROW`] values per row: past that, it runs on the

/// values the rows reference instead.

///

/// `f` can also see values that no row references. If it fails, it runs once

/// more on only the referenced values, so it fails only when running it on

/// each row would.

pub fn map_dictionary_values<F>(array: &ArrayRef, mut f: F) -> Result<ArrayRef>

where

F: FnMut(&ArrayRef) -> Result<ArrayRef>,

{

let compacted = downcast_dictionary_array! {

array => compact_oversized_dictionary(array)?.map(|d| Arc::new(d) as ArrayRef),

dt => return _internal_err!("expected a dictionary array, got {dt}"),

};

let mut dictionary = compacted.unwrap_or_else(|| Arc::clone(array));

let values = match f(dictionary.as_any_dictionary().values()) {

Ok(values) => values,

Err(e) => {

let referenced =

garbage_collect_any_dictionary(dictionary.as_any_dictionary())?;

if referenced.as_any_dictionary().values().len()

== dictionary.as_any_dictionary().values().len()

{

return Err(e);

}

dictionary = referenced;

f(dictionary.as_any_dictionary().values())?

}

};

#[expect(clippy::disallowed_methods)] // the values are bounded above

Ok(dictionary.as_any_dictionary().with_values(values))

}

/// The most dictionary values per row that code processing every value of a

/// dictionary should handle. Past this, it first reduces the dictionary to the

/// values its rows reference. This is a bound, not a tuned break-even: too high

/// only gives up an improvement, while too low can make cheap per-value work

/// slower than processing every value.

pub const MAX_DICTIONARY_VALUES_PER_ROW: usize = 64;

/// Returns `None` when `array` has at most [`MAX_DICTIONARY_VALUES_PER_ROW`]

/// values per row, and otherwise [`compact_dictionary`] of it.

pub(crate) fn compact_oversized_dictionary<K: ArrowDictionaryKeyType>(

array: &DictionaryArray<K>,

) -> Result<Option<DictionaryArray<K>>> {

if array.values().len() <= MAX_DICTIONARY_VALUES_PER_ROW * array.len() {

return Ok(None);

}

compact_dictionary(array).map(Some)

}

/// Returns a copy of `array` that holds each value its rows reference once.

/// This costs O(rows), and work over its values never covers more than the

/// values of `array`, however long they are or often they repeat.

pub fn compact_dictionary<K: ArrowDictionaryKeyType>(

array: &DictionaryArray<K>,

) -> Result<DictionaryArray<K>> {

let keys = array.keys();

let mut new_keys =

HashMap::with_capacity_and_hasher(keys.len(), RandomState::default());

let mut referenced = Vec::with_capacity(keys.len());

let remapped: ScalarBuffer<K::Native> = keys

.iter()

.map(|key| match key {

Some(key) => *new_keys.entry(key.as_usize()).or_insert_with(|| {

referenced.push(key);

// There are at most as many distinct keys as `K` can hold

K::Native::usize_as(referenced.len() - 1)

}),

None => K::Native::default(),

})

.collect();

let values = take(

array.values().as_ref(),

&PrimitiveArray::<K>::from_iter_values(referenced),

None,

)?;

let keys = PrimitiveArray::<K>::new(remapped, keys.nulls().cloned());

Ok(DictionaryArray::try_new(keys, values)?)

}

#[cfg(test)]

mod split_vec_min_alloc_tests {

use super::split_vec_min_alloc;

#[test]

fn drain_branch() {

// n * 2 <= len -> drain+collect branch (allocates n elements)

let mut v = vec![1, 2, 3, 4, 5, 6];

let first = split_vec_min_alloc(&mut v, 2);

assert_eq!(first, vec![1, 2]);

assert_eq!(v, vec![3, 4, 5, 6]);

}

#[test]

fn split_off_branch() {

// remaining < n -> split_off+replace branch (allocates remaining elements)

let mut v = vec![1, 2, 3, 4, 5, 6];

let first = split_vec_min_alloc(&mut v, 4);

assert_eq!(first, vec![1, 2, 3, 4]);

assert_eq!(v, vec![5, 6]);

}

#[test]

fn exactly_half() {

// n * 2 == len -> drain branch (boundary)

let mut v = vec![1, 2, 3, 4];

let first = split_vec_min_alloc(&mut v, 2);

assert_eq!(first, vec![1, 2]);

assert_eq!(v, vec![3, 4]);

}

#[test]

fn take_all() {

let mut v = vec![1, 2, 3];

let first = split_vec_min_alloc(&mut v, 3);

assert_eq!(first, vec![1, 2, 3]);

assert!(v.is_empty());

}

#[test]

fn take_none() {

let mut v = vec![1, 2, 3];

let first = split_vec_min_alloc(&mut v, 0);

assert!(first.is_empty());

assert_eq!(v, vec![1, 2, 3]);

}

#[test]

fn emitted_prefix_does_not_realloc_on_push() {

// Demonstrates *why* the split-off branch must NOT call `shrink_to_fit`.

//

// Downstream callers (e.g. `multi_group_by/bytes.rs`, which does

// `first_n_offsets.push(offset_n)` right after the split) push onto the

// emitted prefix immediately. The split-off branch hands the original

// backing allocation to that prefix, so the prefix already has spare

// capacity for the very next push.

//

// If we shrank the prefix to fit, that next push would have to

// reallocate, and Vec's growth strategy would land it at a *larger*

// capacity than the original allocation we started with -- the opposite

// of the memory saving `shrink_to_fit` was meant to deliver.

// A Vec with a known, deliberately large capacity. n*2 > len, so this

// takes the split-off branch.

let mut v: Vec<u32> = Vec::with_capacity(64);

v.extend(0..10);

let original_capacity = v.capacity();

assert!(original_capacity >= 64);

// Emit a prefix that is most of the Vec (n = 8, remaining = 2).

let mut prefix = split_vec_min_alloc(&mut v, 8);

assert_eq!(prefix, vec![0, 1, 2, 3, 4, 5, 6, 7]);

// The split-off branch moved the original backing store into `prefix`,

// so it keeps the original (large) capacity -- no shrink happened.

assert_eq!(

prefix.capacity(),

original_capacity,

"split-off branch must hand the original allocation to the prefix"

);

// The caller's very next operation: push one element onto the prefix.

prefix.push(99);

// Because the capacity was preserved, the push reused the existing

// allocation: post-push capacity is unchanged and still <= original.

// This is the realloc that `shrink_to_fit` would have forced.

assert_eq!(

prefix.capacity(),

original_capacity,

"push must reuse the preserved allocation (no realloc)"

);

assert!(prefix.capacity() <= original_capacity);

// Counter-demonstration: had we shrunk the prefix to fit (capacity 8),

// the same push would have reallocated. Vec doubles on growth, so the

// post-push capacity (16) ends up LARGER than where a length-8 prefix

// started -- and we paid a realloc for it.

let mut shrunk: Vec<u32> = prefix[..8].to_vec();

shrunk.shrink_to_fit();

let shrunk_capacity = shrink_then_push_capacity(&mut shrunk);

assert!(

shrunk_capacity > 8,

"shrink-to-fit then push reallocates to a larger capacity"

);

}

/// Helper for the counter-demonstration above: push one element and report

/// the resulting capacity.

fn shrink_then_push_capacity(v: &mut Vec<u32>) -> usize {

v.push(99);

v.capacity()

}

}

/// Creates single element [`ListArray`], [`LargeListArray`] and

/// [`FixedSizeListArray`] from other arrays

///

/// For example this builder can convert `[1, 2, 3]` into `[[1, 2, 3]]`

///

/// # Example

/// ```

/// # use std::sync::Arc;

/// # use arrow::array::{Array, ListArray};

/// # use arrow::array::types::Int64Type;

/// # use datafusion_common::utils::SingleRowListArrayBuilder;

/// // Array is [1, 2, 3]

/// let arr = ListArray::from_iter_primitive::<Int64Type, _, _>(vec![Some(vec![

/// Some(1),

/// Some(2),

/// Some(3),

/// ])]);

/// // Wrap as a list array: [[1, 2, 3]]

/// let list_arr = SingleRowListArrayBuilder::new(Arc::new(arr)).build_list_array();

/// assert_eq!(list_arr.len(), 1);

/// ```

#[derive(Debug, Clone)]

pub struct SingleRowListArrayBuilder {

/// array to be wrapped

arr: ArrayRef,

/// Should the resulting array be nullable? Defaults to `true`.

nullable: bool,

/// Specify the field name for the resulting array. Defaults to value used in

/// [`Field::new_list_field`]

field_name: Option<String>,

field_metadata: Metadata,

}

impl SingleRowListArrayBuilder {

/// Create a new instance of [`SingleRowListArrayBuilder`]

pub fn new(arr: ArrayRef) -> Self {

Self {

arr,

nullable: true,

field_name: None,

field_metadata: Metadata::default(),

}

}

/// Set the nullable flag

pub fn with_nullable(mut self, nullable: bool) -> Self {

self.nullable = nullable;

self

}

/// sets the field name for the resulting array

pub fn with_field_name(mut self, field_name: Option<String>) -> Self {

self.field_name = field_name;

self

}

/// Copies field name, nullability, and metadata from the specified field

pub fn with_field(mut self, field: &Field) -> Self {

self.field_name = Some(field.name().to_owned());

self.nullable = field.is_nullable();

self.field_metadata = field.metadata().clone();

self

}

/// Build a single element [`ListArray`]

pub fn build_list_array(self) -> ListArray {

let (field, arr) = self.into_field_and_arr();

let offsets = OffsetBuffer::from_lengths([arr.len()]);

ListArray::new(field, offsets, arr, None)

}

/// Build a single element [`ListArray`] and wrap as [`ScalarValue::List`]

pub fn build_list_scalar(self) -> ScalarValue {

ScalarValue::List(Arc::new(self.build_list_array()))

}

/// Build a single element [`LargeListArray`]

pub fn build_large_list_array(self) -> LargeListArray {

let (field, arr) = self.into_field_and_arr();

let offsets = OffsetBuffer::from_lengths([arr.len()]);

LargeListArray::new(field, offsets, arr, None)

}

/// Build a single element [`LargeListArray`] and wrap as [`ScalarValue::LargeList`]

pub fn build_large_list_scalar(self) -> ScalarValue {

ScalarValue::LargeList(Arc::new(self.build_large_list_array()))

}

/// Build a single element [`FixedSizeListArray`]

pub fn build_fixed_size_list_array(self, list_size: usize) -> FixedSizeListArray {

let (field, arr) = self.into_field_and_arr();

FixedSizeListArray::try_new_with_length(field, list_size as i32, arr, None, 1)

.unwrap()

}

/// Build a single element [`FixedSizeListArray`] and wrap as [`ScalarValue::FixedSizeList`]

pub fn build_fixed_size_list_scalar(self, list_size: usize) -> ScalarValue {

ScalarValue::FixedSizeList(Arc::new(self.build_fixed_size_list_array(list_size)))

}

/// Build a single element [`ListViewArray`]

pub fn build_list_view_array(self) -> ListViewArray {

let (field, arr) = self.into_field_and_arr();

let offsets = ScalarBuffer::from(vec![0]);

let sizes = ScalarBuffer::from(vec![i32::try_from(arr.len()).expect(

"Trying to construct a ListView where element length exceeds i32::MAX",

)]);

ListViewArray::new(field, offsets, sizes, arr, None)

}

/// Build a single element [`ListViewArray`] and wrap as [`ScalarValue::ListView`]

pub fn build_list_view_scalar(self) -> ScalarValue {

ScalarValue::ListView(Arc::new(self.build_list_view_array()))

}

/// Build a single element [`LargeListViewArray`]

pub fn build_large_list_view_array(self) -> LargeListViewArray {

let (field, arr) = self.into_field_and_arr();

let offsets = ScalarBuffer::from(vec![0]);

let sizes = ScalarBuffer::from(vec![arr.len() as i64]);

LargeListViewArray::new(field, offsets, sizes, arr, None)

}

/// Build a single element [`LargeListViewArray`] and wrap as [`ScalarValue::LargeListView`]

pub fn build_large_list_view_scalar(self) -> ScalarValue {

ScalarValue::LargeListView(Arc::new(self.build_large_list_view_array()))

}

/// Helper function: convert this builder into a tuple of field and array

fn into_field_and_arr(self) -> (Arc<Field>, ArrayRef) {

let Self {

arr,

nullable,

field_name,

field_metadata,

} = self;

let data_type = arr.data_type().to_owned();

let field = match field_name {

Some(name) => Field::new(name, data_type, nullable),

None => Field::new_list_field(data_type, nullable),

};

(Arc::new(field.with_metadata(field_metadata)), arr)

}

}

/// Wrap arrays into a single element `ListArray`.

///

/// Example:

/// ```

/// use arrow::array::{Int32Array, ListArray, ArrayRef};

/// use arrow::datatypes::{Int32Type, Field};

/// use std::sync::Arc;

///

/// let arr1 = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef;

/// let arr2 = Arc::new(Int32Array::from(vec![4, 5, 6])) as ArrayRef;

///

/// let list_arr = datafusion_common::utils::arrays_into_list_array([arr1, arr2]).unwrap();

///

/// let expected = ListArray::from_iter_primitive::<Int32Type, _, _>(

/// vec![

/// Some(vec![Some(1), Some(2), Some(3)]),

/// Some(vec![Some(4), Some(5), Some(6)]),

/// ]

/// );

///

/// assert_eq!(list_arr, expected);

/// ```

pub fn arrays_into_list_array(

arr: impl IntoIterator<Item = ArrayRef>,

) -> Result<ListArray> {

let arr = arr.into_iter().collect::<Vec<_>>();

assert_or_internal_err!(!arr.is_empty(), "Cannot wrap empty array into list array");

let lens = arr.iter().map(|x| x.len()).collect::<Vec<_>>();

// Assume data type is consistent

let data_type = arr[0].data_type().to_owned();

let values = arr.iter().map(|x| x.as_ref()).collect::<Vec<_>>();

Ok(ListArray::new(

Arc::new(Field::new_list_field(data_type, true)),

OffsetBuffer::from_lengths(lens),

arrow::compute::concat(values.as_slice())?,

None,

))

}

/// Helper function to convert a ListArray into a vector of ArrayRefs.

pub fn list_to_arrays<O: OffsetSizeTrait>(a: &ArrayRef) -> Vec<ArrayRef> {

a.as_list::<O>().iter().flatten().collect::<Vec<_>>()

}

/// Helper function to convert a FixedSizeListArray into a vector of ArrayRefs.

pub fn fixed_size_list_to_arrays(a: &ArrayRef) -> Vec<ArrayRef> {

a.as_fixed_size_list().iter().flatten().collect::<Vec<_>>()

}

/// Get the base type of a data type.

///

/// Example

/// ```

/// use arrow::datatypes::{DataType, Field};

/// use datafusion_common::utils::base_type;

/// use std::sync::Arc;

///

/// let data_type =

/// DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true)));

/// assert_eq!(base_type(&data_type), DataType::Int32);

///

/// let data_type = DataType::Int32;

/// assert_eq!(base_type(&data_type), DataType::Int32);

/// ```

pub fn base_type(data_type: &DataType) -> DataType {

match data_type {

DataType::List(field)

| DataType::LargeList(field)

| DataType::ListView(field)

| DataType::LargeListView(field)

| DataType::FixedSizeList(field, _) => base_type(field.data_type()),

_ => data_type.to_owned(),

}

}

// TODO: Modify this to also allow specifying how listviews should be treated.

// For example if cast to List (default) or maintain as ListView (requires

// function to implement support for ListViews)

// https://github.com/apache/datafusion/issues/21777

/// Information about how to coerce lists.

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]

pub enum ListCoercion {

/// [`DataType::FixedSizeList`] should be coerced to [`DataType::List`].

FixedSizedListToList,

}

/// A helper function to coerce base type in List.

///

/// Example

/// ```

/// use arrow::datatypes::{DataType, Field};

/// use datafusion_common::utils::coerced_type_with_base_type_only;

/// use std::sync::Arc;

///

/// let data_type = DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true)));

/// let base_type = DataType::Float64;

/// let coerced_type = coerced_type_with_base_type_only(&data_type, &base_type, None);

/// assert_eq!(coerced_type, DataType::List(Arc::new(Field::new_list_field(DataType::Float64, true))));

/// ```

pub fn coerced_type_with_base_type_only(

data_type: &DataType,

base_type: &DataType,

array_coercion: Option<&ListCoercion>,

) -> DataType {

match (data_type, array_coercion) {

(DataType::List(field), _)

| (DataType::FixedSizeList(field, _), Some(ListCoercion::FixedSizedListToList)) =>

{

let field_type = coerced_type_with_base_type_only(

field.data_type(),

base_type,

array_coercion,

);

DataType::List(Arc::new(Field::new(

field.name(),

field_type,

field.is_nullable(),

)))

}

(DataType::FixedSizeList(field, len), _) => {

let field_type = coerced_type_with_base_type_only(

field.data_type(),

base_type,

array_coercion,

);

DataType::FixedSizeList(

Arc::new(Field::new(field.name(), field_type, field.is_nullable())),

*len,

)

}

(DataType::ListView(field), _) => {

let field_type = coerced_type_with_base_type_only(

field.data_type(),

base_type,

array_coercion,

);

DataType::ListView(Arc::new(Field::new(

field.name(),

field_type,

field.is_nullable(),

)))

}

(DataType::LargeList(field), _) => {

let field_type = coerced_type_with_base_type_only(

field.data_type(),

base_type,

array_coercion,

);

DataType::LargeList(Arc::new(Field::new(

field.name(),

field_type,

field.is_nullable(),

)))

}

(DataType::LargeListView(field), _) => {

let field_type = coerced_type_with_base_type_only(

field.data_type(),

base_type,

array_coercion,

);

DataType::LargeListView(Arc::new(Field::new(

field.name(),

field_type,

field.is_nullable(),

)))

}

_ => base_type.clone(),

}

}

/// Recursively coerce and `FixedSizeList` elements to `List`

pub fn coerced_fixed_size_list_to_list(data_type: &DataType) -> DataType {

match data_type {

DataType::List(field) | DataType::FixedSizeList(field, _) => {

let field_type = coerced_fixed_size_list_to_list(field.data_type());

DataType::List(Arc::new(Field::new(

field.name(),

field_type,

field.is_nullable(),

)))

}

DataType::ListView(field) => {

let field_type = coerced_fixed_size_list_to_list(field.data_type());

DataType::ListView(Arc::new(Field::new(

field.name(),

field_type,

field.is_nullable(),

)))

}

DataType::LargeList(field) => {

let field_type = coerced_fixed_size_list_to_list(field.data_type());

DataType::LargeList(Arc::new(Field::new(

field.name(),

field_type,

field.is_nullable(),

)))

}

DataType::LargeListView(field) => {

let field_type = coerced_fixed_size_list_to_list(field.data_type());

DataType::LargeListView(Arc::new(Field::new(

field.name(),

field_type,

field.is_nullable(),

)))

}

_ => data_type.clone(),

}

}

/// Compute the number of dimensions in a list data type.

pub fn list_ndims(data_type: &DataType) -> u64 {

match data_type {

DataType::List(field)

| DataType::LargeList(field)

| DataType::ListView(field)

| DataType::LargeListView(field)

| DataType::FixedSizeList(field, _) => 1 + list_ndims(field.data_type()),

_ => 0,

}

}

/// Adopted from strsim-rs for string similarity metrics

pub mod datafusion_strsim {

// Source: https://github.com/dguo/strsim-rs/blob/master/src/lib.rs

// License: https://github.com/dguo/strsim-rs/blob/master/LICENSE

use std::cmp::min;

use std::str::Chars;

struct StringWrapper<'a>(&'a str);

impl<'b> IntoIterator for &StringWrapper<'b> {

type Item = char;