use arrow2::array::*;
use arrow2::compute::partition::*;
use arrow2::compute::sort::{SortColumn, SortOptions};
use arrow2::datatypes::DataType;
use arrow2::error::Result;
#[test]
fn lexicographical_partition_ranges_empty() {
let input = vec![];
assert!(
lexicographical_partition_ranges(&input).is_err(),
"lexicographical_partition_ranges should reject columns with empty rows"
);
}
#[test]
fn lexicographical_partition_ranges_unaligned_rows() {
let values1 = Int64Array::from([None, Some(-1)]);
let values2 = Utf8Array::<i32>::from([Some("foo")]);
let input = vec![
SortColumn {
values: &values1,
options: None,
},
SortColumn {
values: &values2,
options: None,
},
];
assert!(
lexicographical_partition_ranges(&input).is_err(),
"lexicographical_partition_ranges should reject columns with different row counts"
);
}
#[test]
fn lexicographical_partition_single_column() -> Result<()> {
let values = Int64Array::from_slice([1, 2, 2, 2, 2, 2, 2, 2, 9]);
let input = vec![SortColumn {
values: &values,
options: Some(SortOptions {
descending: false,
nulls_first: true,
}),
}];
{
let results = lexicographical_partition_ranges(&input)?;
assert_eq!(
vec![(0_usize..1_usize), (1_usize..8_usize), (8_usize..9_usize)],
results.collect::<Vec<_>>()
);
}
Ok(())
}
#[test]
fn lexicographical_partition_all_equal_values() -> Result<()> {
let values = Int64Array::from_trusted_len_values_iter(std::iter::repeat(1).take(1000));
let input = vec![SortColumn {
values: &values,
options: Some(SortOptions {
descending: false,
nulls_first: true,
}),
}];
{
let results = lexicographical_partition_ranges(&input)?;
assert_eq!(vec![(0_usize..1000_usize)], results.collect::<Vec<_>>());
}
Ok(())
}
#[test]
fn lexicographical_partition_all_null_values() -> Result<()> {
let values1 = new_null_array(DataType::Int8, 1000);
let values2 = new_null_array(DataType::UInt16, 1000);
let input = vec![
SortColumn {
values: values1.as_ref(),
options: Some(SortOptions {
descending: false,
nulls_first: true,
}),
},
SortColumn {
values: values2.as_ref(),
options: Some(SortOptions {
descending: false,
nulls_first: false,
}),
},
];
{
let results = lexicographical_partition_ranges(&input)?;
assert_eq!(vec![(0_usize..1000_usize)], results.collect::<Vec<_>>());
}
Ok(())
}
#[test]
fn lexicographical_partition_unique_column_1() -> Result<()> {
let values1 = Int64Array::from(vec![None, Some(-1)]);
let values2 = Utf8Array::<i32>::from(vec![Some("foo"), Some("bar")]);
let input = vec![
SortColumn {
values: &values1,
options: Some(SortOptions {
descending: false,
nulls_first: true,
}),
},
SortColumn {
values: &values2,
options: Some(SortOptions {
descending: true,
nulls_first: true,
}),
},
];
{
let results = lexicographical_partition_ranges(&input)?;
assert_eq!(
vec![(0_usize..1_usize), (1_usize..2_usize)],
results.collect::<Vec<_>>()
);
}
Ok(())
}
#[test]
fn lexicographical_partition_unique_column_2() -> Result<()> {
let values1 = Int64Array::from(vec![None, Some(-1), Some(-1)]);
let values2 = Utf8Array::<i32>::from(vec![Some("foo"), Some("bar"), Some("apple")]);
let input = vec![
SortColumn {
values: &values1,
options: Some(SortOptions {
descending: false,
nulls_first: true,
}),
},
SortColumn {
values: &values2,
options: Some(SortOptions {
descending: true,
nulls_first: true,
}),
},
];
{
let results = lexicographical_partition_ranges(&input)?;
assert_eq!(
vec![(0_usize..1_usize), (1_usize..2_usize), (2_usize..3_usize),],
results.collect::<Vec<_>>()
);
}
Ok(())
}
#[test]
fn lexicographical_partition_non_unique_column_1() -> Result<()> {
let values1 = Int64Array::from(vec![None, Some(-1), Some(-1), Some(1)]);
let values2 = Utf8Array::<i32>::from(vec![Some("foo"), Some("bar"), Some("bar"), Some("bar")]);
let input = vec![
SortColumn {
values: &values1,
options: Some(SortOptions {
descending: false,
nulls_first: true,
}),
},
SortColumn {
values: &values2,
options: Some(SortOptions {
descending: true,
nulls_first: true,
}),
},
];
{
let results = lexicographical_partition_ranges(&input)?;
assert_eq!(
vec![(0_usize..1_usize), (1_usize..3_usize), (3_usize..4_usize),],
results.collect::<Vec<_>>()
);
}
Ok(())
}