use std::any::type_name;
use std::sync::Arc;
use arrow::array::*;
use arrow::buffer::OffsetBuffer;
use arrow::compute;
use arrow::datatypes::{DataType, Field, UInt64Type};
use arrow::row::{RowConverter, SortField};
use arrow_buffer::NullBuffer;
use datafusion_common::cast::{
as_generic_string_array, as_int64_array, as_list_array, as_string_array,
};
use datafusion_common::utils::array_into_list_array;
use datafusion_common::{
exec_err, internal_datafusion_err, internal_err, not_impl_err, plan_err,
DataFusionError, Result,
};
use hashbrown::HashSet;
use itertools::Itertools;
macro_rules! downcast_arg {
($ARG:expr, $ARRAY_TYPE:ident) => {{
$ARG.as_any().downcast_ref::<$ARRAY_TYPE>().ok_or_else(|| {
DataFusionError::Internal(format!(
"could not cast to {}",
type_name::<$ARRAY_TYPE>()
))
})?
}};
}
macro_rules! downcast_vec {
($ARGS:expr, $ARRAY_TYPE:ident) => {{
$ARGS
.iter()
.map(|e| match e.as_any().downcast_ref::<$ARRAY_TYPE>() {
Some(array) => Ok(array),
_ => internal_err!("failed to downcast"),
})
}};
}
macro_rules! new_builder {
(BooleanBuilder, $len:expr) => {
BooleanBuilder::with_capacity($len)
};
(StringBuilder, $len:expr) => {
StringBuilder::new()
};
(LargeStringBuilder, $len:expr) => {
LargeStringBuilder::new()
};
($el:ident, $len:expr) => {{
<$el>::with_capacity($len)
}};
}
macro_rules! array {
($ARGS:expr, $ARRAY_TYPE:ident, $BUILDER_TYPE:ident) => {{
let builder = new_builder!($BUILDER_TYPE, $ARGS[0].len());
let mut builder =
ListBuilder::<$BUILDER_TYPE>::with_capacity(builder, $ARGS.len());
let num_rows = $ARGS[0].len();
assert!(
$ARGS.iter().all(|a| a.len() == num_rows),
"all arguments must have the same number of rows"
);
for index in 0..num_rows {
for arg in $ARGS {
match arg.as_any().downcast_ref::<$ARRAY_TYPE>() {
Some(arr) => {
if arr.is_valid(index) {
builder.values().append_value(arr.value(index));
} else {
builder.values().append_null();
}
}
None => match arg.as_any().downcast_ref::<NullArray>() {
Some(arr) => {
for _ in 0..arr.len() {
builder.values().append_null();
}
}
None => return internal_err!("failed to downcast"),
},
}
}
builder.append(true);
}
Arc::new(builder.finish())
}};
}
fn compute_array_length(
arr: Option<ArrayRef>,
dimension: Option<i64>,
) -> Result<Option<u64>> {
let mut current_dimension: i64 = 1;
let mut value = match arr {
Some(arr) => arr,
None => return Ok(None),
};
let dimension = match dimension {
Some(value) => {
if value < 1 {
return Ok(None);
}
value
}
None => return Ok(None),
};
loop {
if current_dimension == dimension {
return Ok(Some(value.len() as u64));
}
match value.data_type() {
DataType::List(..) => {
value = downcast_arg!(value, ListArray).value(0);
current_dimension += 1;
}
_ => return Ok(None),
}
}
}
fn compute_array_ndims(arr: Option<ArrayRef>) -> Result<Option<u64>> {
Ok(compute_array_ndims_with_datatype(arr)?.0)
}
fn compute_array_ndims_with_datatype(
arr: Option<ArrayRef>,
) -> Result<(Option<u64>, DataType)> {
let mut res: u64 = 1;
let mut value = match arr {
Some(arr) => arr,
None => return Ok((None, DataType::Null)),
};
if value.is_empty() {
return Ok((None, DataType::Null));
}
loop {
match value.data_type() {
DataType::List(..) => {
value = downcast_arg!(value, ListArray).value(0);
res += 1;
}
data_type => return Ok((Some(res), data_type.clone())),
}
}
}
fn compute_array_dims(arr: Option<ArrayRef>) -> Result<Option<Vec<Option<u64>>>> {
let mut value = match arr {
Some(arr) => arr,
None => return Ok(None),
};
if value.is_empty() {
return Ok(None);
}
let mut res = vec![Some(value.len() as u64)];
loop {
match value.data_type() {
DataType::List(..) => {
value = downcast_arg!(value, ListArray).value(0);
res.push(Some(value.len() as u64));
}
_ => return Ok(Some(res)),
}
}
}
fn check_datatypes(name: &str, args: &[&ArrayRef]) -> Result<()> {
let data_type = args[0].data_type();
if !args
.iter()
.all(|arg| arg.data_type().equals_datatype(data_type))
{
let types = args.iter().map(|arg| arg.data_type()).collect::<Vec<_>>();
return plan_err!("{name} received incompatible types: '{types:?}'.");
}
Ok(())
}
macro_rules! call_array_function {
($DATATYPE:expr, false) => {
match $DATATYPE {
DataType::Utf8 => array_function!(StringArray),
DataType::LargeUtf8 => array_function!(LargeStringArray),
DataType::Boolean => array_function!(BooleanArray),
DataType::Float32 => array_function!(Float32Array),
DataType::Float64 => array_function!(Float64Array),
DataType::Int8 => array_function!(Int8Array),
DataType::Int16 => array_function!(Int16Array),
DataType::Int32 => array_function!(Int32Array),
DataType::Int64 => array_function!(Int64Array),
DataType::UInt8 => array_function!(UInt8Array),
DataType::UInt16 => array_function!(UInt16Array),
DataType::UInt32 => array_function!(UInt32Array),
DataType::UInt64 => array_function!(UInt64Array),
_ => unreachable!(),
}
};
($DATATYPE:expr, $INCLUDE_LIST:expr) => {{
match $DATATYPE {
DataType::List(_) => array_function!(ListArray),
DataType::Utf8 => array_function!(StringArray),
DataType::LargeUtf8 => array_function!(LargeStringArray),
DataType::Boolean => array_function!(BooleanArray),
DataType::Float32 => array_function!(Float32Array),
DataType::Float64 => array_function!(Float64Array),
DataType::Int8 => array_function!(Int8Array),
DataType::Int16 => array_function!(Int16Array),
DataType::Int32 => array_function!(Int32Array),
DataType::Int64 => array_function!(Int64Array),
DataType::UInt8 => array_function!(UInt8Array),
DataType::UInt16 => array_function!(UInt16Array),
DataType::UInt32 => array_function!(UInt32Array),
DataType::UInt64 => array_function!(UInt64Array),
_ => unreachable!(),
}
}};
}
fn array_array(args: &[ArrayRef], data_type: DataType) -> Result<ArrayRef> {
if args.is_empty() {
return plan_err!("Array requires at least one argument");
}
let res = match data_type {
DataType::List(..) => {
let row_count = args[0].len();
let column_count = args.len();
let mut list_arrays = vec![];
let mut list_array_lengths = vec![];
let mut list_valid = BooleanBufferBuilder::new(row_count);
for index in 0..row_count {
let mut arrays = vec![];
let mut array_lengths = vec![];
let mut valid = BooleanBufferBuilder::new(column_count);
for arg in args {
if arg.as_any().downcast_ref::<NullArray>().is_some() {
array_lengths.push(0);
valid.append(false);
} else {
let list_arr = as_list_array(arg)?;
let arr = list_arr.value(index);
array_lengths.push(arr.len());
arrays.push(arr);
valid.append(true);
}
}
if arrays.is_empty() {
list_valid.append(false);
list_array_lengths.push(0);
} else {
let buffer = valid.finish();
let data_type = arrays[0].data_type();
let field = Arc::new(Field::new("item", data_type.to_owned(), true));
let elements = arrays.iter().map(|x| x.as_ref()).collect::<Vec<_>>();
let values = compute::concat(elements.as_slice())?;
let list_arr = ListArray::new(
field,
OffsetBuffer::from_lengths(array_lengths),
values,
Some(NullBuffer::new(buffer)),
);
list_valid.append(true);
list_array_lengths.push(list_arr.len());
list_arrays.push(list_arr);
}
}
let buffer = list_valid.finish();
let data_type = list_arrays[0].data_type();
let field = Arc::new(Field::new("item", data_type.to_owned(), true));
let elements = list_arrays
.iter()
.map(|x| x as &dyn Array)
.collect::<Vec<_>>();
let values = compute::concat(elements.as_slice())?;
let list_arr = ListArray::new(
field,
OffsetBuffer::from_lengths(list_array_lengths),
values,
Some(NullBuffer::new(buffer)),
);
Arc::new(list_arr)
}
DataType::Utf8 => array!(args, StringArray, StringBuilder),
DataType::LargeUtf8 => array!(args, LargeStringArray, LargeStringBuilder),
DataType::Boolean => array!(args, BooleanArray, BooleanBuilder),
DataType::Float32 => array!(args, Float32Array, Float32Builder),
DataType::Float64 => array!(args, Float64Array, Float64Builder),
DataType::Int8 => array!(args, Int8Array, Int8Builder),
DataType::Int16 => array!(args, Int16Array, Int16Builder),
DataType::Int32 => array!(args, Int32Array, Int32Builder),
DataType::Int64 => array!(args, Int64Array, Int64Builder),
DataType::UInt8 => array!(args, UInt8Array, UInt8Builder),
DataType::UInt16 => array!(args, UInt16Array, UInt16Builder),
DataType::UInt32 => array!(args, UInt32Array, UInt32Builder),
DataType::UInt64 => array!(args, UInt64Array, UInt64Builder),
data_type => {
return not_impl_err!("Array is not implemented for type '{data_type:?}'.")
}
};
Ok(res)
}
pub fn make_array(arrays: &[ArrayRef]) -> Result<ArrayRef> {
let mut data_type = DataType::Null;
for arg in arrays {
let arg_data_type = arg.data_type();
if !arg_data_type.equals_datatype(&DataType::Null) {
data_type = arg_data_type.clone();
break;
}
}
match data_type {
DataType::Null => {
let array = new_null_array(&DataType::Null, arrays.len());
Ok(Arc::new(array_into_list_array(array)))
}
data_type => array_array(arrays, data_type),
}
}
fn return_empty(return_null: bool, data_type: DataType) -> Arc<dyn Array> {
if return_null {
new_null_array(&data_type, 1)
} else {
new_empty_array(&data_type)
}
}
macro_rules! list_slice {
($ARRAY:expr, $I:expr, $J:expr, $RETURN_ELEMENT:expr, $ARRAY_TYPE:ident) => {{
let array = $ARRAY.as_any().downcast_ref::<$ARRAY_TYPE>().unwrap();
if $I == 0 && $J == 0 || $ARRAY.is_empty() {
return return_empty($RETURN_ELEMENT, $ARRAY.data_type().clone());
}
let i = if $I < 0 {
if $I.abs() as usize > array.len() {
return return_empty(true, $ARRAY.data_type().clone());
}
(array.len() as i64 + $I + 1) as usize
} else {
if $I == 0 {
1
} else {
$I as usize
}
};
let j = if $J < 0 {
if $J.abs() as usize > array.len() {
return return_empty(true, $ARRAY.data_type().clone());
}
if $RETURN_ELEMENT {
(array.len() as i64 + $J + 1) as usize
} else {
(array.len() as i64 + $J) as usize
}
} else {
if $J == 0 {
1
} else {
if $J as usize > array.len() {
array.len()
} else {
$J as usize
}
}
};
if i > j || i as usize > $ARRAY.len() {
return_empty($RETURN_ELEMENT, $ARRAY.data_type().clone())
} else {
Arc::new(array.slice((i - 1), (j + 1 - i)))
}
}};
}
macro_rules! slice {
($ARRAY:expr, $KEY:expr, $EXTRA_KEY:expr, $RETURN_ELEMENT:expr, $ARRAY_TYPE:ident) => {{
let sliced_array: Vec<Arc<dyn Array>> = $ARRAY
.iter()
.zip($KEY.iter())
.zip($EXTRA_KEY.iter())
.map(|((arr, i), j)| match (arr, i, j) {
(Some(arr), Some(i), Some(j)) => {
list_slice!(arr, i, j, $RETURN_ELEMENT, $ARRAY_TYPE)
}
(Some(arr), None, Some(j)) => {
list_slice!(arr, 1i64, j, $RETURN_ELEMENT, $ARRAY_TYPE)
}
(Some(arr), Some(i), None) => {
list_slice!(arr, i, arr.len() as i64, $RETURN_ELEMENT, $ARRAY_TYPE)
}
(Some(arr), None, None) if !$RETURN_ELEMENT => arr,
_ => return_empty($RETURN_ELEMENT, $ARRAY.value_type().clone()),
})
.collect();
if sliced_array.is_empty() {
Ok(return_empty($RETURN_ELEMENT, $ARRAY.value_type()))
} else {
let vec = sliced_array
.iter()
.map(|a| a.as_ref())
.collect::<Vec<&dyn Array>>();
let mut i: i32 = 0;
let mut offsets = vec![i];
offsets.extend(
vec.iter()
.map(|a| {
i += a.len() as i32;
i
})
.collect::<Vec<_>>(),
);
let values = compute::concat(vec.as_slice()).unwrap();
if $RETURN_ELEMENT {
Ok(values)
} else {
let field =
Arc::new(Field::new("item", $ARRAY.value_type().clone(), true));
Ok(Arc::new(ListArray::try_new(
field,
OffsetBuffer::new(offsets.into()),
values,
None,
)?))
}
}
}};
}
fn define_array_slice(
list_array: &ListArray,
key: &Int64Array,
extra_key: &Int64Array,
return_element: bool,
) -> Result<ArrayRef> {
macro_rules! array_function {
($ARRAY_TYPE:ident) => {
slice!(list_array, key, extra_key, return_element, $ARRAY_TYPE)
};
}
call_array_function!(list_array.value_type(), true)
}
pub fn array_element(args: &[ArrayRef]) -> Result<ArrayRef> {
let list_array = as_list_array(&args[0])?;
let key = as_int64_array(&args[1])?;
define_array_slice(list_array, key, key, true)
}
pub fn array_slice(args: &[ArrayRef]) -> Result<ArrayRef> {
let list_array = as_list_array(&args[0])?;
let key = as_int64_array(&args[1])?;
let extra_key = as_int64_array(&args[2])?;
define_array_slice(list_array, key, extra_key, false)
}
pub fn array_pop_back(args: &[ArrayRef]) -> Result<ArrayRef> {
let list_array = as_list_array(&args[0])?;
let key = vec![0; list_array.len()];
let extra_key: Vec<_> = list_array
.iter()
.map(|x| x.map_or(0, |arr| arr.len() as i64 - 1))
.collect();
define_array_slice(
list_array,
&Int64Array::from(key),
&Int64Array::from(extra_key),
false,
)
}
pub fn array_append(args: &[ArrayRef]) -> Result<ArrayRef> {
let arr = as_list_array(&args[0])?;
let element = &args[1];
check_datatypes("array_append", &[arr.values(), element])?;
let res = match arr.value_type() {
DataType::List(_) => concat_internal(args)?,
DataType::Null => return make_array(&[element.to_owned()]),
data_type => {
let mut new_values = vec![];
let mut offsets = vec![0];
let elem_data = element.to_data();
for (row_index, arr) in arr.iter().enumerate() {
let new_array = if let Some(arr) = arr {
let original_data = arr.to_data();
let capacity = Capacities::Array(original_data.len() + 1);
let mut mutable = MutableArrayData::with_capacities(
vec![&original_data, &elem_data],
false,
capacity,
);
mutable.extend(0, 0, original_data.len());
mutable.extend(1, row_index, row_index + 1);
let data = mutable.freeze();
arrow_array::make_array(data)
} else {
let capacity = Capacities::Array(1);
let mut mutable = MutableArrayData::with_capacities(
vec![&elem_data],
false,
capacity,
);
mutable.extend(0, row_index, row_index + 1);
let data = mutable.freeze();
arrow_array::make_array(data)
};
offsets.push(offsets[row_index] + new_array.len() as i32);
new_values.push(new_array);
}
let new_values: Vec<_> = new_values.iter().map(|a| a.as_ref()).collect();
let values = arrow::compute::concat(&new_values)?;
Arc::new(ListArray::try_new(
Arc::new(Field::new("item", data_type.to_owned(), true)),
OffsetBuffer::new(offsets.into()),
values,
None,
)?)
}
};
Ok(res)
}
pub fn array_prepend(args: &[ArrayRef]) -> Result<ArrayRef> {
let element = &args[0];
let arr = as_list_array(&args[1])?;
check_datatypes("array_prepend", &[element, arr.values()])?;
let res = match arr.value_type() {
DataType::List(_) => concat_internal(args)?,
DataType::Null => return make_array(&[element.to_owned()]),
data_type => {
let mut new_values = vec![];
let mut offsets = vec![0];
let elem_data = element.to_data();
for (row_index, arr) in arr.iter().enumerate() {
let new_array = if let Some(arr) = arr {
let original_data = arr.to_data();
let capacity = Capacities::Array(original_data.len() + 1);
let mut mutable = MutableArrayData::with_capacities(
vec![&original_data, &elem_data],
false,
capacity,
);
mutable.extend(1, row_index, row_index + 1);
mutable.extend(0, 0, original_data.len());
let data = mutable.freeze();
arrow_array::make_array(data)
} else {
let capacity = Capacities::Array(1);
let mut mutable = MutableArrayData::with_capacities(
vec![&elem_data],
false,
capacity,
);
mutable.extend(0, row_index, row_index + 1);
let data = mutable.freeze();
arrow_array::make_array(data)
};
offsets.push(offsets[row_index] + new_array.len() as i32);
new_values.push(new_array);
}
let new_values: Vec<_> = new_values.iter().map(|a| a.as_ref()).collect();
let values = arrow::compute::concat(&new_values)?;
Arc::new(ListArray::try_new(
Arc::new(Field::new("item", data_type.to_owned(), true)),
OffsetBuffer::new(offsets.into()),
values,
None,
)?)
}
};
Ok(res)
}
fn align_array_dimensions(args: Vec<ArrayRef>) -> Result<Vec<ArrayRef>> {
let args_ndim = args
.iter()
.map(|arg| compute_array_ndims(Some(arg.to_owned())))
.collect::<Result<Vec<_>>>()?
.into_iter()
.map(|x| x.unwrap_or(0))
.collect::<Vec<_>>();
let max_ndim = args_ndim.iter().max().unwrap_or(&0);
let aligned_args: Result<Vec<ArrayRef>> = args
.into_iter()
.zip(args_ndim.iter())
.map(|(array, ndim)| {
if ndim < max_ndim {
let mut aligned_array = array.clone();
for _ in 0..(max_ndim - ndim) {
let data_type = aligned_array.data_type().to_owned();
let array_lengths = vec![1; aligned_array.len()];
let offsets = OffsetBuffer::<i32>::from_lengths(array_lengths);
aligned_array = Arc::new(ListArray::try_new(
Arc::new(Field::new("item", data_type, true)),
offsets,
aligned_array,
None,
)?)
}
Ok(aligned_array)
} else {
Ok(array.clone())
}
})
.collect();
aligned_args
}
fn concat_internal(args: &[ArrayRef]) -> Result<ArrayRef> {
let args = align_array_dimensions(args.to_vec())?;
let list_arrays =
downcast_vec!(args, ListArray).collect::<Result<Vec<&ListArray>>>()?;
let row_count = list_arrays[0].len();
let mut array_lengths = vec![];
let mut arrays = vec![];
let mut valid = BooleanBufferBuilder::new(row_count);
for i in 0..row_count {
let nulls = list_arrays
.iter()
.map(|arr| arr.is_null(i))
.collect::<Vec<_>>();
let is_null = nulls.iter().all(|&x| x);
if is_null {
array_lengths.push(0);
valid.append(false);
} else {
let values = list_arrays
.iter()
.map(|arr| arr.value(i))
.collect::<Vec<_>>();
let elements = values
.iter()
.map(|a| a.as_ref())
.collect::<Vec<&dyn Array>>();
let concated_array = compute::concat(elements.as_slice())?;
array_lengths.push(concated_array.len());
arrays.push(concated_array);
valid.append(true);
}
}
let data_type = list_arrays[0].value_type();
let buffer = valid.finish();
let elements = arrays
.iter()
.map(|a| a.as_ref())
.collect::<Vec<&dyn Array>>();
let list_arr = ListArray::new(
Arc::new(Field::new("item", data_type, true)),
OffsetBuffer::from_lengths(array_lengths),
Arc::new(compute::concat(elements.as_slice())?),
Some(NullBuffer::new(buffer)),
);
Ok(Arc::new(list_arr))
}
pub fn array_concat(args: &[ArrayRef]) -> Result<ArrayRef> {
let mut new_args = vec![];
for arg in args {
let (ndim, lower_data_type) =
compute_array_ndims_with_datatype(Some(arg.clone()))?;
if ndim.is_none() || ndim == Some(1) {
return not_impl_err!("Array is not type '{lower_data_type:?}'.");
} else if !lower_data_type.equals_datatype(&DataType::Null) {
new_args.push(arg.clone());
}
}
concat_internal(new_args.as_slice())
}
pub fn array_empty(args: &[ArrayRef]) -> Result<ArrayRef> {
if args[0].as_any().downcast_ref::<NullArray>().is_some() {
return Ok(Arc::new(BooleanArray::new_null(args[0].len())));
}
let array = as_list_array(&args[0])?;
let builder = array
.iter()
.map(|arr| arr.map(|arr| arr.len() == arr.null_count()))
.collect::<BooleanArray>();
Ok(Arc::new(builder))
}
pub fn array_repeat(args: &[ArrayRef]) -> Result<ArrayRef> {
let element = &args[0];
let count_array = as_int64_array(&args[1])?;
match element.data_type() {
DataType::List(_) => {
let list_array = as_list_array(element)?;
general_list_repeat(list_array, count_array)
}
_ => general_repeat(element, count_array),
}
}
fn general_repeat(array: &ArrayRef, count_array: &Int64Array) -> Result<ArrayRef> {
let data_type = array.data_type();
let mut new_values = vec![];
let count_vec = count_array
.values()
.to_vec()
.iter()
.map(|x| *x as usize)
.collect::<Vec<_>>();
for (row_index, &count) in count_vec.iter().enumerate() {
let repeated_array = if array.is_null(row_index) {
new_null_array(data_type, count)
} else {
let original_data = array.to_data();
let capacity = Capacities::Array(count);
let mut mutable =
MutableArrayData::with_capacities(vec![&original_data], false, capacity);
for _ in 0..count {
mutable.extend(0, row_index, row_index + 1);
}
let data = mutable.freeze();
arrow_array::make_array(data)
};
new_values.push(repeated_array);
}
let new_values: Vec<_> = new_values.iter().map(|a| a.as_ref()).collect();
let values = compute::concat(&new_values)?;
Ok(Arc::new(ListArray::try_new(
Arc::new(Field::new("item", data_type.to_owned(), true)),
OffsetBuffer::from_lengths(count_vec),
values,
None,
)?))
}
fn general_list_repeat(
list_array: &ListArray,
count_array: &Int64Array,
) -> Result<ArrayRef> {
let data_type = list_array.data_type();
let value_type = list_array.value_type();
let mut new_values = vec![];
let count_vec = count_array
.values()
.to_vec()
.iter()
.map(|x| *x as usize)
.collect::<Vec<_>>();
for (list_array_row, &count) in list_array.iter().zip(count_vec.iter()) {
let list_arr = match list_array_row {
Some(list_array_row) => {
let original_data = list_array_row.to_data();
let capacity = Capacities::Array(original_data.len() * count);
let mut mutable = MutableArrayData::with_capacities(
vec![&original_data],
false,
capacity,
);
for _ in 0..count {
mutable.extend(0, 0, original_data.len());
}
let data = mutable.freeze();
let repeated_array = arrow_array::make_array(data);
let list_arr = ListArray::try_new(
Arc::new(Field::new("item", value_type.clone(), true)),
OffsetBuffer::from_lengths(vec![original_data.len(); count]),
repeated_array,
None,
)?;
Arc::new(list_arr) as ArrayRef
}
None => new_null_array(data_type, count),
};
new_values.push(list_arr);
}
let lengths = new_values.iter().map(|a| a.len()).collect::<Vec<_>>();
let new_values: Vec<_> = new_values.iter().map(|a| a.as_ref()).collect();
let values = compute::concat(&new_values)?;
Ok(Arc::new(ListArray::try_new(
Arc::new(Field::new("item", data_type.to_owned(), true)),
OffsetBuffer::from_lengths(lengths),
values,
None,
)?))
}
macro_rules! position {
($ARRAY:expr, $ELEMENT:expr, $INDEX:expr, $ARRAY_TYPE:ident) => {{
let element = downcast_arg!($ELEMENT, $ARRAY_TYPE);
$ARRAY
.iter()
.zip(element.iter())
.zip($INDEX.iter())
.map(|((arr, el), i)| {
let index = match i {
Some(i) => {
if i <= 0 {
0
} else {
i - 1
}
}
None => return exec_err!("initial position must not be null"),
};
match arr {
Some(arr) => {
let child_array = downcast_arg!(arr, $ARRAY_TYPE);
match child_array
.iter()
.skip(index as usize)
.position(|x| x == el)
{
Some(value) => Ok(Some(value as u64 + index as u64 + 1u64)),
None => Ok(None),
}
}
None => Ok(None),
}
})
.collect::<Result<UInt64Array>>()?
}};
}
pub fn array_position(args: &[ArrayRef]) -> Result<ArrayRef> {
let arr = as_list_array(&args[0])?;
let element = &args[1];
let index = if args.len() == 3 {
as_int64_array(&args[2])?.clone()
} else {
Int64Array::from_value(0, arr.len())
};
check_datatypes("array_position", &[arr.values(), element])?;
macro_rules! array_function {
($ARRAY_TYPE:ident) => {
position!(arr, element, index, $ARRAY_TYPE)
};
}
let res = call_array_function!(arr.value_type(), true);
Ok(Arc::new(res))
}
macro_rules! positions {
($ARRAY:expr, $ELEMENT:expr, $ARRAY_TYPE:ident) => {{
let element = downcast_arg!($ELEMENT, $ARRAY_TYPE);
let mut offsets: Vec<i32> = vec![0];
let mut values =
downcast_arg!(new_empty_array(&DataType::UInt64), UInt64Array).clone();
for comp in $ARRAY
.iter()
.zip(element.iter())
.map(|(arr, el)| match arr {
Some(arr) => {
let child_array = downcast_arg!(arr, $ARRAY_TYPE);
let res = child_array
.iter()
.enumerate()
.filter(|(_, x)| *x == el)
.flat_map(|(i, _)| Some((i + 1) as u64))
.collect::<UInt64Array>();
Ok(res)
}
None => Ok(downcast_arg!(
new_empty_array(&DataType::UInt64),
UInt64Array
)
.clone()),
})
.collect::<Result<Vec<UInt64Array>>>()?
{
let last_offset: i32 = offsets.last().copied().ok_or_else(|| {
DataFusionError::Internal(format!("offsets should not be empty",))
})?;
values =
downcast_arg!(compute::concat(&[&values, &comp,])?.clone(), UInt64Array)
.clone();
offsets.push(last_offset + comp.len() as i32);
}
let field = Arc::new(Field::new("item", DataType::UInt64, true));
Arc::new(ListArray::try_new(
field,
OffsetBuffer::new(offsets.into()),
Arc::new(values),
None,
)?)
}};
}
pub fn array_positions(args: &[ArrayRef]) -> Result<ArrayRef> {
let arr = as_list_array(&args[0])?;
let element = &args[1];
check_datatypes("array_positions", &[arr.values(), element])?;
macro_rules! array_function {
($ARRAY_TYPE:ident) => {
positions!(arr, element, $ARRAY_TYPE)
};
}
let res = call_array_function!(arr.value_type(), true);
Ok(res)
}
macro_rules! general_remove {
($ARRAY:expr, $ELEMENT:expr, $MAX:expr, $ARRAY_TYPE:ident) => {{
let mut offsets: Vec<i32> = vec![0];
let mut values =
downcast_arg!(new_empty_array($ELEMENT.data_type()), $ARRAY_TYPE).clone();
let element = downcast_arg!($ELEMENT, $ARRAY_TYPE);
for ((arr, el), max) in $ARRAY.iter().zip(element.iter()).zip($MAX.iter()) {
let last_offset: i32 = offsets.last().copied().ok_or_else(|| {
DataFusionError::Internal(format!("offsets should not be empty"))
})?;
match arr {
Some(arr) => {
let child_array = downcast_arg!(arr, $ARRAY_TYPE);
let mut counter = 0;
let max = if max < Some(1) { 1 } else { max.unwrap() };
let filter_array = child_array
.iter()
.map(|element| {
if counter != max && element == el {
counter += 1;
Some(false)
} else {
Some(true)
}
})
.collect::<BooleanArray>();
let filtered_array = compute::filter(&child_array, &filter_array)?;
values = downcast_arg!(
compute::concat(&[&values, &filtered_array,])?.clone(),
$ARRAY_TYPE
)
.clone();
offsets.push(last_offset + filtered_array.len() as i32);
}
None => offsets.push(last_offset),
}
}
let field = Arc::new(Field::new("item", $ELEMENT.data_type().clone(), true));
Arc::new(ListArray::try_new(
field,
OffsetBuffer::new(offsets.into()),
Arc::new(values),
None,
)?)
}};
}
macro_rules! array_removement_function {
($FUNC:ident, $MAX_FUNC:expr, $DOC:expr) => {
#[doc = $DOC]
pub fn $FUNC(args: &[ArrayRef]) -> Result<ArrayRef> {
let arr = as_list_array(&args[0])?;
let element = &args[1];
let max = $MAX_FUNC(args)?;
check_datatypes(stringify!($FUNC), &[arr.values(), element])?;
macro_rules! array_function {
($ARRAY_TYPE:ident) => {
general_remove!(arr, element, max, $ARRAY_TYPE)
};
}
let res = call_array_function!(arr.value_type(), true);
Ok(res)
}
};
}
fn remove_one(args: &[ArrayRef]) -> Result<Int64Array> {
Ok(Int64Array::from_value(1, args[0].len()))
}
fn remove_n(args: &[ArrayRef]) -> Result<Int64Array> {
as_int64_array(&args[2]).cloned()
}
fn remove_all(args: &[ArrayRef]) -> Result<Int64Array> {
Ok(Int64Array::from_value(i64::MAX, args[0].len()))
}
array_removement_function!(array_remove, remove_one, "Array_remove SQL function");
array_removement_function!(array_remove_n, remove_n, "Array_remove_n SQL function");
array_removement_function!(
array_remove_all,
remove_all,
"Array_remove_all SQL function"
);
fn general_replace(
list_array: &ListArray,
from_array: &ArrayRef,
to_array: &ArrayRef,
arr_n: Vec<i64>,
) -> Result<ArrayRef> {
let mut offsets: Vec<i32> = vec![0];
let data_type = list_array.value_type();
let mut new_values = vec![];
for (row_index, (list_array_row, n)) in
list_array.iter().zip(arr_n.iter()).enumerate()
{
let last_offset: i32 = offsets
.last()
.copied()
.ok_or_else(|| internal_datafusion_err!("offsets should not be empty"))?;
match list_array_row {
Some(list_array_row) => {
let indices = UInt32Array::from(vec![row_index as u32]);
let from_array_row = arrow::compute::take(from_array, &indices, None)?;
let eq_array = match from_array_row.data_type() {
DataType::List(_) => {
let from_array_row_inner =
as_list_array(&from_array_row)?.value(0);
let list_array_row_inner = as_list_array(&list_array_row)?;
list_array_row_inner
.iter()
.map(|row| row.map(|row| row.eq(&from_array_row_inner)))
.collect::<BooleanArray>()
}
_ => {
let from_arr = Scalar::new(from_array_row);
arrow_ord::cmp::not_distinct(&list_array_row, &from_arr)?
}
};
let original_data = list_array_row.to_data();
let to_data = to_array.to_data();
let capacity = Capacities::Array(original_data.len() + to_data.len());
let mut mutable = MutableArrayData::with_capacities(
vec![&original_data, &to_data],
false,
capacity,
);
let original_idx = 0;
let replace_idx = 1;
let mut counter = 0;
for (i, to_replace) in eq_array.iter().enumerate() {
if let Some(true) = to_replace {
mutable.extend(replace_idx, row_index, row_index + 1);
counter += 1;
if counter == *n {
mutable.extend(original_idx, i + 1, eq_array.len());
break;
}
} else {
mutable.extend(original_idx, i, i + 1);
}
}
let data = mutable.freeze();
let replaced_array = arrow_array::make_array(data);
offsets.push(last_offset + replaced_array.len() as i32);
new_values.push(replaced_array);
}
None => {
offsets.push(last_offset);
}
}
}
let values = if new_values.is_empty() {
new_empty_array(&data_type)
} else {
let new_values: Vec<_> = new_values.iter().map(|a| a.as_ref()).collect();
arrow::compute::concat(&new_values)?
};
Ok(Arc::new(ListArray::try_new(
Arc::new(Field::new("item", data_type, true)),
OffsetBuffer::new(offsets.into()),
values,
list_array.nulls().cloned(),
)?))
}
pub fn array_replace(args: &[ArrayRef]) -> Result<ArrayRef> {
let arr_n = vec![1; args[0].len()];
general_replace(as_list_array(&args[0])?, &args[1], &args[2], arr_n)
}
pub fn array_replace_n(args: &[ArrayRef]) -> Result<ArrayRef> {
let arr_n = as_int64_array(&args[3])?.values().to_vec();
general_replace(as_list_array(&args[0])?, &args[1], &args[2], arr_n)
}
pub fn array_replace_all(args: &[ArrayRef]) -> Result<ArrayRef> {
let arr_n = vec![i64::MAX; args[0].len()];
general_replace(as_list_array(&args[0])?, &args[1], &args[2], arr_n)
}
macro_rules! to_string {
($ARG:expr, $ARRAY:expr, $DELIMITER:expr, $NULL_STRING:expr, $WITH_NULL_STRING:expr, $ARRAY_TYPE:ident) => {{
let arr = downcast_arg!($ARRAY, $ARRAY_TYPE);
for x in arr {
match x {
Some(x) => {
$ARG.push_str(&x.to_string());
$ARG.push_str($DELIMITER);
}
None => {
if $WITH_NULL_STRING {
$ARG.push_str($NULL_STRING);
$ARG.push_str($DELIMITER);
}
}
}
}
Ok($ARG)
}};
}
pub fn array_to_string(args: &[ArrayRef]) -> Result<ArrayRef> {
let arr = &args[0];
let delimiters = as_string_array(&args[1])?;
let delimiters: Vec<Option<&str>> = delimiters.iter().collect();
let mut null_string = String::from("");
let mut with_null_string = false;
if args.len() == 3 {
null_string = as_string_array(&args[2])?.value(0).to_string();
with_null_string = true;
}
fn compute_array_to_string(
arg: &mut String,
arr: ArrayRef,
delimiter: String,
null_string: String,
with_null_string: bool,
) -> Result<&mut String> {
match arr.data_type() {
DataType::List(..) => {
let list_array = downcast_arg!(arr, ListArray);
for i in 0..list_array.len() {
compute_array_to_string(
arg,
list_array.value(i),
delimiter.clone(),
null_string.clone(),
with_null_string,
)?;
}
Ok(arg)
}
DataType::Null => Ok(arg),
data_type => {
macro_rules! array_function {
($ARRAY_TYPE:ident) => {
to_string!(
arg,
arr,
&delimiter,
&null_string,
with_null_string,
$ARRAY_TYPE
)
};
}
call_array_function!(data_type, false)
}
}
}
let mut arg = String::from("");
let mut res: Vec<Option<String>> = Vec::new();
match arr.data_type() {
DataType::List(_) | DataType::LargeList(_) | DataType::FixedSizeList(_, _) => {
let list_array = arr.as_list::<i32>();
for (arr, &delimiter) in list_array.iter().zip(delimiters.iter()) {
if let (Some(arr), Some(delimiter)) = (arr, delimiter) {
arg = String::from("");
let s = compute_array_to_string(
&mut arg,
arr,
delimiter.to_string(),
null_string.clone(),
with_null_string,
)?
.clone();
if let Some(s) = s.strip_suffix(delimiter) {
res.push(Some(s.to_string()));
} else {
res.push(Some(s));
}
} else {
res.push(None);
}
}
}
_ => {
assert_eq!(delimiters.len(), 1);
let delimiter = delimiters[0].unwrap();
let s = compute_array_to_string(
&mut arg,
arr.clone(),
delimiter.to_string(),
null_string,
with_null_string,
)?
.clone();
if !s.is_empty() {
let s = s.strip_suffix(delimiter).unwrap().to_string();
res.push(Some(s));
} else {
res.push(Some(s));
}
}
}
Ok(Arc::new(StringArray::from(res)))
}
pub fn cardinality(args: &[ArrayRef]) -> Result<ArrayRef> {
let list_array = as_list_array(&args[0])?.clone();
let result = list_array
.iter()
.map(|arr| match compute_array_dims(arr)? {
Some(vector) => Ok(Some(vector.iter().map(|x| x.unwrap()).product::<u64>())),
None => Ok(None),
})
.collect::<Result<UInt64Array>>()?;
Ok(Arc::new(result) as ArrayRef)
}
fn get_offsets_for_flatten(
offsets: OffsetBuffer<i32>,
indexes: OffsetBuffer<i32>,
) -> OffsetBuffer<i32> {
let buffer = offsets.into_inner();
let offsets: Vec<i32> = indexes.iter().map(|i| buffer[*i as usize]).collect();
OffsetBuffer::new(offsets.into())
}
fn flatten_internal(
array: &dyn Array,
indexes: Option<OffsetBuffer<i32>>,
) -> Result<ListArray> {
let list_arr = as_list_array(array)?;
let (field, offsets, values, _) = list_arr.clone().into_parts();
let data_type = field.data_type();
match data_type {
DataType::List(_) => {
if let Some(indexes) = indexes {
let offsets = get_offsets_for_flatten(offsets, indexes);
flatten_internal(&values, Some(offsets))
} else {
flatten_internal(&values, Some(offsets))
}
}
_ => {
if let Some(indexes) = indexes {
let offsets = get_offsets_for_flatten(offsets, indexes);
let list_arr = ListArray::new(field, offsets, values, None);
Ok(list_arr)
} else {
Ok(list_arr.clone())
}
}
}
}
pub fn flatten(args: &[ArrayRef]) -> Result<ArrayRef> {
let flattened_array = flatten_internal(&args[0], None)?;
Ok(Arc::new(flattened_array) as ArrayRef)
}
pub fn array_length(args: &[ArrayRef]) -> Result<ArrayRef> {
let list_array = as_list_array(&args[0])?;
let dimension = if args.len() == 2 {
as_int64_array(&args[1])?.clone()
} else {
Int64Array::from_value(1, list_array.len())
};
let result = list_array
.iter()
.zip(dimension.iter())
.map(|(arr, dim)| compute_array_length(arr, dim))
.collect::<Result<UInt64Array>>()?;
Ok(Arc::new(result) as ArrayRef)
}
pub fn array_dims(args: &[ArrayRef]) -> Result<ArrayRef> {
let list_array = as_list_array(&args[0])?;
let data = list_array
.iter()
.map(compute_array_dims)
.collect::<Result<Vec<_>>>()?;
let result = ListArray::from_iter_primitive::<UInt64Type, _, _>(data);
Ok(Arc::new(result) as ArrayRef)
}
pub fn array_ndims(args: &[ArrayRef]) -> Result<ArrayRef> {
let list_array = as_list_array(&args[0])?;
let result = list_array
.iter()
.map(compute_array_ndims)
.collect::<Result<UInt64Array>>()?;
Ok(Arc::new(result) as ArrayRef)
}
macro_rules! non_list_contains {
($ARRAY:expr, $SUB_ARRAY:expr, $ARRAY_TYPE:ident) => {{
let sub_array = downcast_arg!($SUB_ARRAY, $ARRAY_TYPE);
let mut boolean_builder = BooleanArray::builder($ARRAY.len());
for (arr, elem) in $ARRAY.iter().zip(sub_array.iter()) {
if let (Some(arr), Some(elem)) = (arr, elem) {
let arr = downcast_arg!(arr, $ARRAY_TYPE);
let res = arr.iter().dedup().flatten().any(|x| x == elem);
boolean_builder.append_value(res);
}
}
Ok(Arc::new(boolean_builder.finish()))
}};
}
pub fn array_has(args: &[ArrayRef]) -> Result<ArrayRef> {
let array = as_list_array(&args[0])?;
let element = &args[1];
check_datatypes("array_has", &[array.values(), element])?;
match element.data_type() {
DataType::List(_) => {
let sub_array = as_list_array(element)?;
let mut boolean_builder = BooleanArray::builder(array.len());
for (arr, elem) in array.iter().zip(sub_array.iter()) {
if let (Some(arr), Some(elem)) = (arr, elem) {
let list_arr = as_list_array(&arr)?;
let res = list_arr.iter().dedup().flatten().any(|x| *x == *elem);
boolean_builder.append_value(res);
}
}
Ok(Arc::new(boolean_builder.finish()))
}
data_type => {
macro_rules! array_function {
($ARRAY_TYPE:ident) => {
non_list_contains!(array, element, $ARRAY_TYPE)
};
}
call_array_function!(data_type, false)
}
}
}
macro_rules! array_has_any_non_list_check {
($ARRAY:expr, $SUB_ARRAY:expr, $ARRAY_TYPE:ident) => {{
let arr = downcast_arg!($ARRAY, $ARRAY_TYPE);
let sub_arr = downcast_arg!($SUB_ARRAY, $ARRAY_TYPE);
let mut res = false;
for elem in sub_arr.iter().dedup() {
if let Some(elem) = elem {
res |= arr.iter().dedup().flatten().any(|x| x == elem);
} else {
return internal_err!(
"array_has_any does not support Null type for element in sub_array"
);
}
}
res
}};
}
pub fn array_has_any(args: &[ArrayRef]) -> Result<ArrayRef> {
check_datatypes("array_has_any", &[&args[0], &args[1]])?;
let array = as_list_array(&args[0])?;
let sub_array = as_list_array(&args[1])?;
let mut boolean_builder = BooleanArray::builder(array.len());
for (arr, sub_arr) in array.iter().zip(sub_array.iter()) {
if let (Some(arr), Some(sub_arr)) = (arr, sub_arr) {
let res = match arr.data_type() {
DataType::List(_) => {
let arr = downcast_arg!(arr, ListArray);
let sub_arr = downcast_arg!(sub_arr, ListArray);
let mut res = false;
for elem in sub_arr.iter().dedup().flatten() {
res |= arr.iter().dedup().flatten().any(|x| *x == *elem);
}
res
}
data_type => {
macro_rules! array_function {
($ARRAY_TYPE:ident) => {
array_has_any_non_list_check!(arr, sub_arr, $ARRAY_TYPE)
};
}
call_array_function!(data_type, false)
}
};
boolean_builder.append_value(res);
}
}
Ok(Arc::new(boolean_builder.finish()))
}
macro_rules! array_has_all_non_list_check {
($ARRAY:expr, $SUB_ARRAY:expr, $ARRAY_TYPE:ident) => {{
let arr = downcast_arg!($ARRAY, $ARRAY_TYPE);
let sub_arr = downcast_arg!($SUB_ARRAY, $ARRAY_TYPE);
let mut res = true;
for elem in sub_arr.iter().dedup() {
if let Some(elem) = elem {
res &= arr.iter().dedup().flatten().any(|x| x == elem);
} else {
return internal_err!(
"array_has_all does not support Null type for element in sub_array"
);
}
}
res
}};
}
pub fn array_has_all(args: &[ArrayRef]) -> Result<ArrayRef> {
check_datatypes("array_has_all", &[&args[0], &args[1]])?;
let array = as_list_array(&args[0])?;
let sub_array = as_list_array(&args[1])?;
let mut boolean_builder = BooleanArray::builder(array.len());
for (arr, sub_arr) in array.iter().zip(sub_array.iter()) {
if let (Some(arr), Some(sub_arr)) = (arr, sub_arr) {
let res = match arr.data_type() {
DataType::List(_) => {
let arr = downcast_arg!(arr, ListArray);
let sub_arr = downcast_arg!(sub_arr, ListArray);
let mut res = true;
for elem in sub_arr.iter().dedup().flatten() {
res &= arr.iter().dedup().flatten().any(|x| *x == *elem);
}
res
}
data_type => {
macro_rules! array_function {
($ARRAY_TYPE:ident) => {
array_has_all_non_list_check!(arr, sub_arr, $ARRAY_TYPE)
};
}
call_array_function!(data_type, false)
}
};
boolean_builder.append_value(res);
}
}
Ok(Arc::new(boolean_builder.finish()))
}
pub fn string_to_array<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> {
let string_array = as_generic_string_array::<T>(&args[0])?;
let delimiter_array = as_generic_string_array::<T>(&args[1])?;
let mut list_builder = ListBuilder::new(StringBuilder::with_capacity(
string_array.len(),
string_array.get_buffer_memory_size(),
));
match args.len() {
2 => {
string_array.iter().zip(delimiter_array.iter()).for_each(
|(string, delimiter)| {
match (string, delimiter) {
(Some(string), Some("")) => {
list_builder.values().append_value(string);
list_builder.append(true);
}
(Some(string), Some(delimiter)) => {
string.split(delimiter).for_each(|s| {
list_builder.values().append_value(s);
});
list_builder.append(true);
}
(Some(string), None) => {
string.chars().map(|c| c.to_string()).for_each(|c| {
list_builder.values().append_value(c);
});
list_builder.append(true);
}
_ => list_builder.append(false), }
},
);
}
3 => {
let null_value_array = as_generic_string_array::<T>(&args[2])?;
string_array
.iter()
.zip(delimiter_array.iter())
.zip(null_value_array.iter())
.for_each(|((string, delimiter), null_value)| {
match (string, delimiter) {
(Some(string), Some("")) => {
if Some(string) == null_value {
list_builder.values().append_null();
} else {
list_builder.values().append_value(string);
}
list_builder.append(true);
}
(Some(string), Some(delimiter)) => {
string.split(delimiter).for_each(|s| {
if Some(s) == null_value {
list_builder.values().append_null();
} else {
list_builder.values().append_value(s);
}
});
list_builder.append(true);
}
(Some(string), None) => {
string.chars().map(|c| c.to_string()).for_each(|c| {
if Some(c.as_str()) == null_value {
list_builder.values().append_null();
} else {
list_builder.values().append_value(c);
}
});
list_builder.append(true);
}
_ => list_builder.append(false), }
});
}
_ => {
return internal_err!(
"Expect string_to_array function to take two or three parameters"
)
}
}
let list_array = list_builder.finish();
Ok(Arc::new(list_array) as ArrayRef)
}
pub fn array_intersect(args: &[ArrayRef]) -> Result<ArrayRef> {
assert_eq!(args.len(), 2);
let first_array = as_list_array(&args[0])?;
let second_array = as_list_array(&args[1])?;
if first_array.value_type() != second_array.value_type() {
return internal_err!("array_intersect is not implemented for '{first_array:?}' and '{second_array:?}'");
}
let dt = first_array.value_type().clone();
let mut offsets = vec![0];
let mut new_arrays = vec![];
let converter = RowConverter::new(vec![SortField::new(dt.clone())])?;
for (first_arr, second_arr) in first_array.iter().zip(second_array.iter()) {
if let (Some(first_arr), Some(second_arr)) = (first_arr, second_arr) {
let l_values = converter.convert_columns(&[first_arr])?;
let r_values = converter.convert_columns(&[second_arr])?;
let values_set: HashSet<_> = l_values.iter().collect();
let mut rows = Vec::with_capacity(r_values.num_rows());
for r_val in r_values.iter().sorted().dedup() {
if values_set.contains(&r_val) {
rows.push(r_val);
}
}
let last_offset: i32 = match offsets.last().copied() {
Some(offset) => offset,
None => return internal_err!("offsets should not be empty"),
};
offsets.push(last_offset + rows.len() as i32);
let arrays = converter.convert_rows(rows)?;
let array = match arrays.get(0) {
Some(array) => array.clone(),
None => {
return internal_err!(
"array_intersect: failed to get array from rows"
)
}
};
new_arrays.push(array);
}
}
let field = Arc::new(Field::new("item", dt, true));
let offsets = OffsetBuffer::new(offsets.into());
let new_arrays_ref = new_arrays.iter().map(|v| v.as_ref()).collect::<Vec<_>>();
let values = compute::concat(&new_arrays_ref)?;
let arr = Arc::new(ListArray::try_new(field, offsets, values, None)?);
Ok(arr)
}
#[cfg(test)]
mod tests {
use super::*;
use arrow::datatypes::Int64Type;
use datafusion_common::cast::as_uint64_array;
#[test]
fn test_align_array_dimensions() {
let array1d_1 =
Arc::new(ListArray::from_iter_primitive::<Int64Type, _, _>(vec![
Some(vec![Some(1), Some(2), Some(3)]),
Some(vec![Some(4), Some(5)]),
]));
let array1d_2 =
Arc::new(ListArray::from_iter_primitive::<Int64Type, _, _>(vec![
Some(vec![Some(6), Some(7), Some(8)]),
]));
let array2d_1 = Arc::new(array_into_list_array(array1d_1.clone())) as ArrayRef;
let array2d_2 = Arc::new(array_into_list_array(array1d_2.clone())) as ArrayRef;
let res =
align_array_dimensions(vec![array1d_1.to_owned(), array2d_2.to_owned()])
.unwrap();
let expected = as_list_array(&array2d_1).unwrap();
let expected_dim = compute_array_ndims(Some(array2d_1.to_owned())).unwrap();
assert_ne!(as_list_array(&res[0]).unwrap(), expected);
assert_eq!(
compute_array_ndims(Some(res[0].clone())).unwrap(),
expected_dim
);
let array3d_1 = Arc::new(array_into_list_array(array2d_1)) as ArrayRef;
let array3d_2 = array_into_list_array(array2d_2.to_owned());
let res =
align_array_dimensions(vec![array1d_1, Arc::new(array3d_2.clone())]).unwrap();
let expected = as_list_array(&array3d_1).unwrap();
let expected_dim = compute_array_ndims(Some(array3d_1.to_owned())).unwrap();
assert_ne!(as_list_array(&res[0]).unwrap(), expected);
assert_eq!(
compute_array_ndims(Some(res[0].clone())).unwrap(),
expected_dim
);
}
#[test]
fn test_array() {
let args = [
Arc::new(Int64Array::from(vec![1])) as ArrayRef,
Arc::new(Int64Array::from(vec![2])),
Arc::new(Int64Array::from(vec![3])),
];
let array = make_array(&args).expect("failed to initialize function array");
let result = as_list_array(&array).expect("failed to initialize function array");
assert_eq!(result.len(), 1);
assert_eq!(
&[1, 2, 3],
as_int64_array(&result.value(0))
.expect("failed to cast to primitive array")
.values()
)
}
#[test]
fn test_nested_array() {
let args = [
Arc::new(Int64Array::from(vec![1, 2])) as ArrayRef,
Arc::new(Int64Array::from(vec![3, 4])),
Arc::new(Int64Array::from(vec![5, 6])),
];
let array = make_array(&args).expect("failed to initialize function array");
let result = as_list_array(&array).expect("failed to initialize function array");
assert_eq!(result.len(), 2);
assert_eq!(
&[1, 3, 5],
as_int64_array(&result.value(0))
.expect("failed to cast to primitive array")
.values()
);
assert_eq!(
&[2, 4, 6],
as_int64_array(&result.value(1))
.expect("failed to cast to primitive array")
.values()
);
}
#[test]
fn test_array_element() {
let list_array = return_array();
let arr = array_element(&[list_array, Arc::new(Int64Array::from_value(1, 1))])
.expect("failed to initialize function array_element");
let result =
as_int64_array(&arr).expect("failed to initialize function array_element");
assert_eq!(result, &Int64Array::from_value(1, 1));
let list_array = return_array();
let arr = array_element(&[list_array, Arc::new(Int64Array::from_value(3, 1))])
.expect("failed to initialize function array_element");
let result =
as_int64_array(&arr).expect("failed to initialize function array_element");
assert_eq!(result, &Int64Array::from_value(3, 1));
let list_array = return_array();
let arr = array_element(&[list_array, Arc::new(Int64Array::from_value(0, 1))])
.expect("failed to initialize function array_element");
let result =
as_int64_array(&arr).expect("failed to initialize function array_element");
assert_eq!(result, &Int64Array::from(vec![None]));
let list_array = return_array();
let arr = array_element(&[list_array, Arc::new(Int64Array::from(vec![None]))])
.expect("failed to initialize function array_element");
let result =
as_int64_array(&arr).expect("failed to initialize function array_element");
assert_eq!(result, &Int64Array::from(vec![None]));
let list_array = return_array();
let arr = array_element(&[list_array, Arc::new(Int64Array::from_value(-1, 1))])
.expect("failed to initialize function array_element");
let result =
as_int64_array(&arr).expect("failed to initialize function array_element");
assert_eq!(result, &Int64Array::from_value(4, 1));
let list_array = return_array();
let arr = array_element(&[list_array, Arc::new(Int64Array::from_value(-3, 1))])
.expect("failed to initialize function array_element");
let result =
as_int64_array(&arr).expect("failed to initialize function array_element");
assert_eq!(result, &Int64Array::from_value(2, 1));
let list_array = return_array();
let arr = array_element(&[list_array, Arc::new(Int64Array::from_value(10, 1))])
.expect("failed to initialize function array_element");
let result =
as_int64_array(&arr).expect("failed to initialize function array_element");
assert_eq!(result, &Int64Array::from(vec![None]));
}
#[test]
fn test_nested_array_element() {
let list_array = return_nested_array();
let arr = array_element(&[list_array, Arc::new(Int64Array::from_value(2, 1))])
.expect("failed to initialize function array_element");
let result =
as_list_array(&arr).expect("failed to initialize function array_element");
assert_eq!(
&[5, 6, 7, 8],
result
.value(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
);
}
#[test]
fn test_array_pop_back() {
let list_array = return_array();
let arr = array_pop_back(&[list_array])
.expect("failed to initialize function array_pop_back");
let result =
as_list_array(&arr).expect("failed to initialize function array_pop_back");
assert_eq!(
&[1, 2, 3],
result
.value(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
);
let list_array = Arc::new(result.clone());
let arr = array_pop_back(&[list_array])
.expect("failed to initialize function array_pop_back");
let result =
as_list_array(&arr).expect("failed to initialize function array_pop_back");
assert_eq!(
&[1, 2],
result
.value(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
);
let list_array = Arc::new(result.clone());
let arr = array_pop_back(&[list_array])
.expect("failed to initialize function array_pop_back");
let result =
as_list_array(&arr).expect("failed to initialize function array_pop_back");
assert_eq!(
&[1],
result
.value(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
);
let list_array = Arc::new(result.clone());
let arr = array_pop_back(&[list_array])
.expect("failed to initialize function array_pop_back");
let result =
as_list_array(&arr).expect("failed to initialize function array_pop_back");
assert_eq!(
&[],
result
.value(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
);
let list_array = Arc::new(result.clone());
let arr = array_pop_back(&[list_array])
.expect("failed to initialize function array_pop_back");
let result =
as_list_array(&arr).expect("failed to initialize function array_pop_back");
assert_eq!(
&[],
result
.value(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
);
let list_array = return_array_with_nulls();
let arr = array_pop_back(&[list_array])
.expect("failed to initialize function array_pop_back");
let result =
as_list_array(&arr).expect("failed to initialize function array_pop_back");
assert_eq!(3, result.values().len());
assert_eq!(
&[false, true, false],
&[
result.values().is_null(0),
result.values().is_null(1),
result.values().is_null(2)
]
);
}
#[test]
fn test_nested_array_pop_back() {
let list_array = return_nested_array();
let arr = array_pop_back(&[list_array])
.expect("failed to initialize function array_slice");
let result =
as_list_array(&arr).expect("failed to initialize function array_slice");
assert_eq!(
&[1, 2, 3, 4],
result
.value(0)
.as_any()
.downcast_ref::<ListArray>()
.unwrap()
.value(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
);
let list_array = Arc::new(result.clone());
let arr = array_pop_back(&[list_array])
.expect("failed to initialize function array_pop_back");
let result =
as_list_array(&arr).expect("failed to initialize function array_pop_back");
assert!(result
.value(0)
.as_any()
.downcast_ref::<ListArray>()
.unwrap()
.is_empty());
let list_array = Arc::new(result.clone());
let arr = array_pop_back(&[list_array])
.expect("failed to initialize function array_pop_back");
let result =
as_list_array(&arr).expect("failed to initialize function array_pop_back");
assert!(result
.value(0)
.as_any()
.downcast_ref::<ListArray>()
.unwrap()
.is_empty());
}
#[test]
fn test_array_slice() {
let list_array = return_array();
let arr = array_slice(&[
list_array,
Arc::new(Int64Array::from_value(1, 1)),
Arc::new(Int64Array::from_value(3, 1)),
])
.expect("failed to initialize function array_slice");
let result =
as_list_array(&arr).expect("failed to initialize function array_slice");
assert_eq!(
&[1, 2, 3],
result
.value(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
);
let list_array = return_array();
let arr = array_slice(&[
list_array,
Arc::new(Int64Array::from_value(2, 1)),
Arc::new(Int64Array::from_value(2, 1)),
])
.expect("failed to initialize function array_slice");
let result =
as_list_array(&arr).expect("failed to initialize function array_slice");
assert_eq!(
&[2],
result
.value(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
);
let list_array = return_array();
let arr = array_slice(&[
list_array,
Arc::new(Int64Array::from_value(0, 1)),
Arc::new(Int64Array::from_value(0, 1)),
])
.expect("failed to initialize function array_slice");
let result =
as_list_array(&arr).expect("failed to initialize function array_slice");
assert!(result
.value(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.is_empty());
let list_array = return_array();
let arr = array_slice(&[
list_array,
Arc::new(Int64Array::from_value(0, 1)),
Arc::new(Int64Array::from_value(6, 1)),
])
.expect("failed to initialize function array_slice");
let result =
as_list_array(&arr).expect("failed to initialize function array_slice");
assert_eq!(
&[1, 2, 3, 4],
result
.value(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
);
let list_array = return_array();
let arr = array_slice(&[
list_array,
Arc::new(Int64Array::from_value(-2, 1)),
Arc::new(Int64Array::from_value(-2, 1)),
])
.expect("failed to initialize function array_slice");
let result =
as_list_array(&arr).expect("failed to initialize function array_slice");
assert!(result
.value(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.is_empty());
let list_array = return_array();
let arr = array_slice(&[
list_array,
Arc::new(Int64Array::from_value(-3, 1)),
Arc::new(Int64Array::from_value(-1, 1)),
])
.expect("failed to initialize function array_slice");
let result =
as_list_array(&arr).expect("failed to initialize function array_slice");
assert_eq!(
&[2, 3],
result
.value(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
);
let list_array = return_array();
let arr = array_slice(&[
list_array,
Arc::new(Int64Array::from_value(-3, 1)),
Arc::new(Int64Array::from_value(2, 1)),
])
.expect("failed to initialize function array_slice");
let result =
as_list_array(&arr).expect("failed to initialize function array_slice");
assert_eq!(
&[2],
result
.value(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
);
let list_array = return_array();
let arr = array_slice(&[
list_array,
Arc::new(Int64Array::from_value(2, 1)),
Arc::new(Int64Array::from_value(11, 1)),
])
.expect("failed to initialize function array_slice");
let result =
as_list_array(&arr).expect("failed to initialize function array_slice");
assert_eq!(
&[2, 3, 4],
result
.value(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
);
let list_array = return_array();
let arr = array_slice(&[
list_array,
Arc::new(Int64Array::from_value(3, 1)),
Arc::new(Int64Array::from_value(1, 1)),
])
.expect("failed to initialize function array_slice");
let result =
as_list_array(&arr).expect("failed to initialize function array_slice");
assert!(result
.value(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.is_empty());
let list_array = return_array();
let arr = array_slice(&[
list_array,
Arc::new(Int64Array::from_value(-7, 1)),
Arc::new(Int64Array::from_value(-2, 1)),
])
.expect("failed to initialize function array_slice");
let result =
as_list_array(&arr).expect("failed to initialize function array_slice");
assert!(result
.value(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.is_null(0));
}
#[test]
fn test_nested_array_slice() {
let list_array = return_nested_array();
let arr = array_slice(&[
list_array,
Arc::new(Int64Array::from_value(1, 1)),
Arc::new(Int64Array::from_value(1, 1)),
])
.expect("failed to initialize function array_slice");
let result =
as_list_array(&arr).expect("failed to initialize function array_slice");
assert_eq!(
&[1, 2, 3, 4],
result
.value(0)
.as_any()
.downcast_ref::<ListArray>()
.unwrap()
.value(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
);
let list_array = return_nested_array();
let arr = array_slice(&[
list_array,
Arc::new(Int64Array::from_value(-1, 1)),
Arc::new(Int64Array::from_value(-1, 1)),
])
.expect("failed to initialize function array_slice");
let result =
as_list_array(&arr).expect("failed to initialize function array_slice");
assert!(result
.value(0)
.as_any()
.downcast_ref::<ListArray>()
.unwrap()
.is_empty());
let list_array = return_nested_array();
let arr = array_slice(&[
list_array,
Arc::new(Int64Array::from_value(-1, 1)),
Arc::new(Int64Array::from_value(2, 1)),
])
.expect("failed to initialize function array_slice");
let result =
as_list_array(&arr).expect("failed to initialize function array_slice");
assert_eq!(
&[5, 6, 7, 8],
result
.value(0)
.as_any()
.downcast_ref::<ListArray>()
.unwrap()
.value(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
);
}
#[test]
fn test_array_append() {
let data = vec![Some(vec![Some(1), Some(2), Some(3)])];
let list_array =
Arc::new(ListArray::from_iter_primitive::<Int64Type, _, _>(data)) as ArrayRef;
let int64_array = Arc::new(Int64Array::from(vec![Some(4)])) as ArrayRef;
let args = [list_array, int64_array];
let array =
array_append(&args).expect("failed to initialize function array_append");
let result =
as_list_array(&array).expect("failed to initialize function array_append");
assert_eq!(
&[1, 2, 3, 4],
result
.value(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
);
}
#[test]
fn test_array_prepend() {
let data = vec![Some(vec![Some(2), Some(3), Some(4)])];
let list_array =
Arc::new(ListArray::from_iter_primitive::<Int64Type, _, _>(data)) as ArrayRef;
let int64_array = Arc::new(Int64Array::from(vec![Some(1)])) as ArrayRef;
let args = [int64_array, list_array];
let array =
array_prepend(&args).expect("failed to initialize function array_append");
let result =
as_list_array(&array).expect("failed to initialize function array_append");
assert_eq!(
&[1, 2, 3, 4],
result
.value(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
);
}
#[test]
fn test_array_concat() {
let data = vec![Some(vec![Some(1), Some(2), Some(3)])];
let list_array1 =
Arc::new(ListArray::from_iter_primitive::<Int64Type, _, _>(data)) as ArrayRef;
let data = vec![Some(vec![Some(4), Some(5), Some(6)])];
let list_array2 =
Arc::new(ListArray::from_iter_primitive::<Int64Type, _, _>(data)) as ArrayRef;
let data = vec![Some(vec![Some(7), Some(8), Some(9)])];
let list_array3 =
Arc::new(ListArray::from_iter_primitive::<Int64Type, _, _>(data)) as ArrayRef;
let args = [list_array1, list_array2, list_array3];
let array =
array_concat(&args).expect("failed to initialize function array_concat");
let result =
as_list_array(&array).expect("failed to initialize function array_concat");
assert_eq!(
&[1, 2, 3, 4, 5, 6, 7, 8, 9],
result
.value(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
);
}
#[test]
fn test_nested_array_concat() {
let list_array = return_array();
let arr = array_concat(&[list_array.clone(), list_array.clone()])
.expect("failed to initialize function array_concat");
let result =
as_list_array(&arr).expect("failed to initialize function array_concat");
assert_eq!(
&[1, 2, 3, 4, 1, 2, 3, 4],
result
.value(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
);
let list_nested_array = return_nested_array();
let list_array = return_array();
let arr = array_concat(&[list_nested_array, list_array])
.expect("failed to initialize function array_concat");
let result =
as_list_array(&arr).expect("failed to initialize function array_concat");
assert_eq!(
&[1, 2, 3, 4],
result
.value(0)
.as_any()
.downcast_ref::<ListArray>()
.unwrap()
.value(2)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
);
}
#[test]
fn test_array_position() {
let list_array = return_array();
let array = array_position(&[list_array, Arc::new(Int64Array::from_value(3, 1))])
.expect("failed to initialize function array_position");
let result = as_uint64_array(&array)
.expect("failed to initialize function array_position");
assert_eq!(result, &UInt64Array::from(vec![3]));
}
#[test]
fn test_array_positions() {
let list_array = return_array();
let array =
array_positions(&[list_array, Arc::new(Int64Array::from_value(3, 1))])
.expect("failed to initialize function array_position");
let result =
as_list_array(&array).expect("failed to initialize function array_position");
assert_eq!(result.len(), 1);
assert_eq!(
&[3],
result
.value(0)
.as_any()
.downcast_ref::<UInt64Array>()
.unwrap()
.values()
);
}
#[test]
fn test_array_remove() {
let list_array = return_array_with_repeating_elements();
let array = array_remove(&[list_array, Arc::new(Int64Array::from_value(3, 1))])
.expect("failed to initialize function array_remove");
let result =
as_list_array(&array).expect("failed to initialize function array_remove");
assert_eq!(result.len(), 1);
assert_eq!(
&[1, 2, 3, 2, 3],
result
.value(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
);
}
#[test]
fn test_nested_array_remove() {
let list_array = return_nested_array_with_repeating_elements();
let element_array = return_array();
let array = array_remove(&[list_array, element_array])
.expect("failed to initialize function array_remove");
let result =
as_list_array(&array).expect("failed to initialize function array_remove");
assert_eq!(result.len(), 1);
let data = vec![
Some(vec![Some(5), Some(6), Some(7), Some(8)]),
Some(vec![Some(1), Some(2), Some(3), Some(4)]),
Some(vec![Some(9), Some(10), Some(11), Some(12)]),
Some(vec![Some(5), Some(6), Some(7), Some(8)]),
];
let expected = ListArray::from_iter_primitive::<Int64Type, _, _>(data);
assert_eq!(
expected,
result
.value(0)
.as_any()
.downcast_ref::<ListArray>()
.unwrap()
.clone()
);
}
#[test]
fn test_array_remove_n() {
let list_array = return_array_with_repeating_elements();
let array = array_remove_n(&[
list_array,
Arc::new(Int64Array::from_value(3, 1)),
Arc::new(Int64Array::from_value(2, 1)),
])
.expect("failed to initialize function array_remove_n");
let result =
as_list_array(&array).expect("failed to initialize function array_remove_n");
assert_eq!(result.len(), 1);
assert_eq!(
&[1, 2, 2, 3],
result
.value(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
);
}
#[test]
fn test_nested_array_remove_n() {
let list_array = return_nested_array_with_repeating_elements();
let element_array = return_array();
let array = array_remove_n(&[
list_array,
element_array,
Arc::new(Int64Array::from_value(3, 1)),
])
.expect("failed to initialize function array_remove_n");
let result =
as_list_array(&array).expect("failed to initialize function array_remove_n");
assert_eq!(result.len(), 1);
let data = vec![
Some(vec![Some(5), Some(6), Some(7), Some(8)]),
Some(vec![Some(9), Some(10), Some(11), Some(12)]),
Some(vec![Some(5), Some(6), Some(7), Some(8)]),
];
let expected = ListArray::from_iter_primitive::<Int64Type, _, _>(data);
assert_eq!(
expected,
result
.value(0)
.as_any()
.downcast_ref::<ListArray>()
.unwrap()
.clone()
);
}
#[test]
fn test_array_remove_all() {
let list_array = return_array_with_repeating_elements();
let array =
array_remove_all(&[list_array, Arc::new(Int64Array::from_value(3, 1))])
.expect("failed to initialize function array_remove_all");
let result = as_list_array(&array)
.expect("failed to initialize function array_remove_all");
assert_eq!(result.len(), 1);
assert_eq!(
&[1, 2, 2],
result
.value(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
);
}
#[test]
fn test_nested_array_remove_all() {
let list_array = return_nested_array_with_repeating_elements();
let element_array = return_array();
let array = array_remove_all(&[list_array, element_array])
.expect("failed to initialize function array_remove_all");
let result = as_list_array(&array)
.expect("failed to initialize function array_remove_all");
assert_eq!(result.len(), 1);
let data = vec![
Some(vec![Some(5), Some(6), Some(7), Some(8)]),
Some(vec![Some(9), Some(10), Some(11), Some(12)]),
Some(vec![Some(5), Some(6), Some(7), Some(8)]),
];
let expected = ListArray::from_iter_primitive::<Int64Type, _, _>(data);
assert_eq!(
expected,
result
.value(0)
.as_any()
.downcast_ref::<ListArray>()
.unwrap()
.clone()
);
}
#[test]
fn test_array_to_string() {
let list_array = return_array();
let array =
array_to_string(&[list_array, Arc::new(StringArray::from(vec![Some(",")]))])
.expect("failed to initialize function array_to_string");
let result = as_string_array(&array)
.expect("failed to initialize function array_to_string");
assert_eq!(result.len(), 1);
assert_eq!("1,2,3,4", result.value(0));
let list_array = return_array_with_nulls();
let array = array_to_string(&[
list_array,
Arc::new(StringArray::from(vec![Some(",")])),
Arc::new(StringArray::from(vec![Some("*")])),
])
.expect("failed to initialize function array_to_string");
let result = as_string_array(&array)
.expect("failed to initialize function array_to_string");
assert_eq!(result.len(), 1);
assert_eq!("1,*,3,*", result.value(0));
}
#[test]
fn test_nested_array_to_string() {
let list_array = return_nested_array();
let array =
array_to_string(&[list_array, Arc::new(StringArray::from(vec![Some("-")]))])
.expect("failed to initialize function array_to_string");
let result = as_string_array(&array)
.expect("failed to initialize function array_to_string");
assert_eq!(result.len(), 1);
assert_eq!("1-2-3-4-5-6-7-8", result.value(0));
let list_array = return_nested_array_with_nulls();
let array = array_to_string(&[
list_array,
Arc::new(StringArray::from(vec![Some("-")])),
Arc::new(StringArray::from(vec![Some("*")])),
])
.expect("failed to initialize function array_to_string");
let result = as_string_array(&array)
.expect("failed to initialize function array_to_string");
assert_eq!(result.len(), 1);
assert_eq!("1-*-3-*-*-6-7-*", result.value(0));
}
#[test]
fn test_cardinality() {
let list_array = return_array();
let arr = cardinality(&[list_array])
.expect("failed to initialize function cardinality");
let result =
as_uint64_array(&arr).expect("failed to initialize function cardinality");
assert_eq!(result, &UInt64Array::from(vec![4]));
}
#[test]
fn test_nested_cardinality() {
let list_array = return_nested_array();
let arr = cardinality(&[list_array])
.expect("failed to initialize function cardinality");
let result =
as_uint64_array(&arr).expect("failed to initialize function cardinality");
assert_eq!(result, &UInt64Array::from(vec![8]));
}
#[test]
fn test_array_length() {
let list_array = return_array();
let arr = array_length(&[list_array.clone()])
.expect("failed to initialize function array_ndims");
let result =
as_uint64_array(&arr).expect("failed to initialize function array_ndims");
assert_eq!(result, &UInt64Array::from_value(4, 1));
let array = array_length(&[list_array, Arc::new(Int64Array::from_value(1, 1))])
.expect("failed to initialize function array_ndims");
let result =
as_uint64_array(&array).expect("failed to initialize function array_ndims");
assert_eq!(result, &UInt64Array::from_value(4, 1));
}
#[test]
fn test_nested_array_length() {
let list_array = return_nested_array();
let arr = array_length(&[list_array.clone()])
.expect("failed to initialize function array_length");
let result =
as_uint64_array(&arr).expect("failed to initialize function array_length");
assert_eq!(result, &UInt64Array::from_value(2, 1));
let arr =
array_length(&[list_array.clone(), Arc::new(Int64Array::from_value(1, 1))])
.expect("failed to initialize function array_length");
let result =
as_uint64_array(&arr).expect("failed to initialize function array_length");
assert_eq!(result, &UInt64Array::from_value(2, 1));
let arr =
array_length(&[list_array.clone(), Arc::new(Int64Array::from_value(2, 1))])
.expect("failed to initialize function array_length");
let result =
as_uint64_array(&arr).expect("failed to initialize function array_length");
assert_eq!(result, &UInt64Array::from_value(4, 1));
let arr = array_length(&[list_array, Arc::new(Int64Array::from_value(3, 1))])
.expect("failed to initialize function array_length");
let result =
as_uint64_array(&arr).expect("failed to initialize function array_length");
assert_eq!(result, &UInt64Array::from(vec![None]));
}
#[test]
fn test_array_dims() {
let list_array = return_array();
let array =
array_dims(&[list_array]).expect("failed to initialize function array_dims");
let result =
as_list_array(&array).expect("failed to initialize function array_dims");
assert_eq!(
&[4],
result
.value(0)
.as_any()
.downcast_ref::<UInt64Array>()
.unwrap()
.values()
);
}
#[test]
fn test_nested_array_dims() {
let list_array = return_nested_array();
let array =
array_dims(&[list_array]).expect("failed to initialize function array_dims");
let result =
as_list_array(&array).expect("failed to initialize function array_dims");
assert_eq!(
&[2, 4],
result
.value(0)
.as_any()
.downcast_ref::<UInt64Array>()
.unwrap()
.values()
);
}
#[test]
fn test_array_ndims() {
let list_array = return_array();
let array = array_ndims(&[list_array])
.expect("failed to initialize function array_ndims");
let result =
as_uint64_array(&array).expect("failed to initialize function array_ndims");
assert_eq!(result, &UInt64Array::from_value(1, 1));
}
#[test]
fn test_nested_array_ndims() {
let list_array = return_nested_array();
let array = array_ndims(&[list_array])
.expect("failed to initialize function array_ndims");
let result =
as_uint64_array(&array).expect("failed to initialize function array_ndims");
assert_eq!(result, &UInt64Array::from_value(2, 1));
}
#[test]
fn test_check_invalid_datatypes() {
let data = vec![Some(vec![Some(1), Some(2), Some(3)])];
let list_array =
Arc::new(ListArray::from_iter_primitive::<Int64Type, _, _>(data)) as ArrayRef;
let int64_array = Arc::new(StringArray::from(vec![Some("string")])) as ArrayRef;
let args = [list_array.clone(), int64_array.clone()];
let array = array_append(&args);
assert_eq!(array.unwrap_err().strip_backtrace(), "Error during planning: array_append received incompatible types: '[Int64, Utf8]'.");
}
fn return_array() -> ArrayRef {
let args = [
Arc::new(Int64Array::from(vec![Some(1)])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(2)])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(3)])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(4)])) as ArrayRef,
];
make_array(&args).expect("failed to initialize function array")
}
fn return_nested_array() -> ArrayRef {
let args = [
Arc::new(Int64Array::from(vec![Some(1)])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(2)])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(3)])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(4)])) as ArrayRef,
];
let arr1 = make_array(&args).expect("failed to initialize function array");
let args = [
Arc::new(Int64Array::from(vec![Some(5)])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(6)])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(7)])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(8)])) as ArrayRef,
];
let arr2 = make_array(&args).expect("failed to initialize function array");
make_array(&[arr1, arr2]).expect("failed to initialize function array")
}
fn return_array_with_nulls() -> ArrayRef {
let args = [
Arc::new(Int64Array::from(vec![Some(1)])) as ArrayRef,
Arc::new(Int64Array::from(vec![None])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(3)])) as ArrayRef,
Arc::new(Int64Array::from(vec![None])) as ArrayRef,
];
make_array(&args).expect("failed to initialize function array")
}
fn return_nested_array_with_nulls() -> ArrayRef {
let args = [
Arc::new(Int64Array::from(vec![Some(1)])) as ArrayRef,
Arc::new(Int64Array::from(vec![None])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(3)])) as ArrayRef,
Arc::new(Int64Array::from(vec![None])) as ArrayRef,
];
let arr1 = make_array(&args).expect("failed to initialize function array");
let args = [
Arc::new(Int64Array::from(vec![None])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(6)])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(7)])) as ArrayRef,
Arc::new(Int64Array::from(vec![None])) as ArrayRef,
];
let arr2 = make_array(&args).expect("failed to initialize function array");
make_array(&[arr1, arr2]).expect("failed to initialize function array")
}
fn return_array_with_repeating_elements() -> ArrayRef {
let args = [
Arc::new(Int64Array::from(vec![Some(3)])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(1)])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(2)])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(3)])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(2)])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(3)])) as ArrayRef,
];
make_array(&args).expect("failed to initialize function array")
}
fn return_nested_array_with_repeating_elements() -> ArrayRef {
let args = [
Arc::new(Int64Array::from(vec![Some(1)])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(2)])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(3)])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(4)])) as ArrayRef,
];
let arr1 = make_array(&args).expect("failed to initialize function array");
let args = [
Arc::new(Int64Array::from(vec![Some(5)])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(6)])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(7)])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(8)])) as ArrayRef,
];
let arr2 = make_array(&args).expect("failed to initialize function array");
let args = [
Arc::new(Int64Array::from(vec![Some(1)])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(2)])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(3)])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(4)])) as ArrayRef,
];
let arr3 = make_array(&args).expect("failed to initialize function array");
let args = [
Arc::new(Int64Array::from(vec![Some(9)])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(10)])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(11)])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(12)])) as ArrayRef,
];
let arr4 = make_array(&args).expect("failed to initialize function array");
let args = [
Arc::new(Int64Array::from(vec![Some(5)])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(6)])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(7)])) as ArrayRef,
Arc::new(Int64Array::from(vec![Some(8)])) as ArrayRef,
];
let arr5 = make_array(&args).expect("failed to initialize function array");
make_array(&[arr1, arr2, arr3, arr4, arr5])
.expect("failed to initialize function array")
}
}