use std::sync::Arc;
use vortex_buffer::BitBuffer;
use vortex_error::VortexExpect;
use vortex_error::vortex_panic;
use vortex_mask::Mask;
use crate::ArrayRef;
use crate::IntoArray;
use crate::LEGACY_SESSION;
use crate::VortexSessionExecute;
use crate::arrays::BoolArray;
use crate::arrays::PrimitiveArray;
use crate::builtins::ArrayBuiltins;
use crate::dtype::DType;
use crate::dtype::Nullability;
use crate::dtype::PType;
use crate::scalar_fn::fns::operators::Operator;
fn test_filter_take_consistency(array: &ArrayRef) {
let len = array.len();
if len == 0 {
return;
}
let mask_pattern: BitBuffer = (0..len).map(|i| i % 3 != 1).collect();
let mask = Mask::from_buffer(mask_pattern.clone());
let filtered = array
.filter(mask)
.vortex_expect("filter should succeed in conformance test");
let indices: Vec<u64> = mask_pattern
.iter()
.enumerate()
.filter_map(|(i, v)| v.then_some(i as u64))
.collect();
let indices_array = PrimitiveArray::from_iter(indices).into_array();
let taken = array
.take(indices_array)
.vortex_expect("take should succeed in conformance test");
assert_eq!(
filtered.len(),
taken.len(),
"Filter and take should produce arrays of the same length. \
Filtered length: {}, Taken length: {}",
filtered.len(),
taken.len()
);
for i in 0..filtered.len() {
let filtered_val = filtered
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test");
let taken_val = taken
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test");
assert_eq!(
filtered_val, taken_val,
"Filter and take produced different values at index {i}. \
Filtered value: {filtered_val:?}, Taken value: {taken_val:?}"
);
}
}
fn test_double_mask_consistency(array: &ArrayRef) {
let len = array.len();
if len == 0 {
return;
}
let mask1: Mask = (0..len).map(|i| i % 3 == 0).collect();
let mask2: Mask = (0..len).map(|i| i % 2 == 0).collect();
let first_masked = array
.clone()
.mask((!&mask1).into_array())
.vortex_expect("mask should succeed in conformance test");
let double_masked = first_masked
.mask((!&mask2).into_array())
.vortex_expect("mask should succeed in conformance test");
let combined_pattern: BitBuffer = mask1
.to_bit_buffer()
.iter()
.zip(mask2.to_bit_buffer().iter())
.map(|(a, b)| a || b)
.collect();
let combined_mask = Mask::from_buffer(combined_pattern);
let directly_masked = array
.clone()
.mask((!&combined_mask).into_array())
.vortex_expect("mask should succeed in conformance test");
assert_eq!(
double_masked.len(),
directly_masked.len(),
"Sequential masking and combined masking should produce arrays of the same length. \
Sequential length: {}, Combined length: {}",
double_masked.len(),
directly_masked.len()
);
for i in 0..double_masked.len() {
let double_val = double_masked
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test");
let direct_val = directly_masked
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test");
assert_eq!(
double_val, direct_val,
"Sequential masking and combined masking produced different values at index {i}. \
Sequential masking value: {double_val:?}, Combined masking value: {direct_val:?}\n\
This likely indicates an issue with how masks are composed in the array implementation."
);
}
}
fn test_filter_identity(array: &ArrayRef) {
let len = array.len();
if len == 0 {
return;
}
let all_true_mask = Mask::new_true(len);
let filtered = array
.filter(all_true_mask)
.vortex_expect("filter should succeed in conformance test");
assert_eq!(
filtered.len(),
array.len(),
"Filtering with all-true mask should preserve array length. \
Original length: {}, Filtered length: {}",
array.len(),
filtered.len()
);
for i in 0..len {
let original_val = array
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test");
let filtered_val = filtered
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test");
assert_eq!(
filtered_val, original_val,
"Filtering with all-true mask should preserve all values. \
Value at index {i} changed from {original_val:?} to {filtered_val:?}"
);
}
}
fn test_mask_identity(array: &ArrayRef) {
let len = array.len();
if len == 0 {
return;
}
let all_false_mask = Mask::new_false(len);
let masked = array
.clone()
.mask((!&all_false_mask).into_array())
.vortex_expect("mask should succeed in conformance test");
assert_eq!(
masked.len(),
array.len(),
"Masking with all-false mask should preserve array length. \
Original length: {}, Masked length: {}",
array.len(),
masked.len()
);
assert!(
masked.dtype().is_nullable(),
"Mask operation should always produce a nullable array, but dtype is {}",
masked.dtype()
);
for i in 0..len {
let original_val = array
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test");
let masked_val = masked
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test");
let expected_val = original_val.clone().into_nullable();
assert_eq!(
masked_val, expected_val,
"Masking with all-false mask should preserve values (as nullable). \
Value at index {i}: original = {original_val:?}, masked = {masked_val:?}, expected = {expected_val:?}"
);
}
}
fn test_slice_filter_consistency(array: &ArrayRef) {
let len = array.len();
if len < 4 {
return; }
let mut mask_pattern = vec![false; len];
mask_pattern[1..4.min(len)].fill(true);
let mask = Mask::from_iter(mask_pattern);
let filtered = array
.filter(mask)
.vortex_expect("filter should succeed in conformance test");
let sliced = array
.slice(1..4.min(len))
.vortex_expect("slice should succeed in conformance test");
assert_eq!(
filtered.len(),
sliced.len(),
"Filter with contiguous mask and slice should produce same length. \
Filtered length: {}, Sliced length: {}",
filtered.len(),
sliced.len()
);
for i in 0..filtered.len() {
let filtered_val = filtered
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test");
let sliced_val = sliced
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test");
assert_eq!(
filtered_val, sliced_val,
"Filter with contiguous mask and slice produced different values at index {i}. \
Filtered value: {filtered_val:?}, Sliced value: {sliced_val:?}"
);
}
}
fn test_take_slice_consistency(array: &ArrayRef) {
let len = array.len();
if len < 3 {
return; }
let end = 4.min(len);
let indices = PrimitiveArray::from_iter((1..end).map(|i| i as u64)).into_array();
let taken = array
.take(indices)
.vortex_expect("take should succeed in conformance test");
let sliced = array
.slice(1..end)
.vortex_expect("slice should succeed in conformance test");
assert_eq!(
taken.len(),
sliced.len(),
"Take with sequential indices and slice should produce same length. \
Taken length: {}, Sliced length: {}",
taken.len(),
sliced.len()
);
for i in 0..taken.len() {
let taken_val = taken
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test");
let sliced_val = sliced
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test");
assert_eq!(
taken_val, sliced_val,
"Take with sequential indices and slice produced different values at index {i}. \
Taken value: {taken_val:?}, Sliced value: {sliced_val:?}"
);
}
}
fn test_filter_preserves_order(array: &ArrayRef) {
let len = array.len();
if len < 4 {
return;
}
let mask_pattern: Vec<bool> = (0..len).map(|i| i == 0 || i == 2 || i == 3).collect();
let mask = Mask::from_iter(mask_pattern);
let filtered = array
.filter(mask)
.vortex_expect("filter should succeed in conformance test");
assert_eq!(filtered.len(), 3.min(len));
if len >= 4 {
assert_eq!(
filtered
.scalar_at(0)
.vortex_expect("scalar_at should succeed in conformance test"),
array
.scalar_at(0)
.vortex_expect("scalar_at should succeed in conformance test")
);
assert_eq!(
filtered
.scalar_at(1)
.vortex_expect("scalar_at should succeed in conformance test"),
array
.scalar_at(2)
.vortex_expect("scalar_at should succeed in conformance test")
);
assert_eq!(
filtered
.scalar_at(2)
.vortex_expect("scalar_at should succeed in conformance test"),
array
.scalar_at(3)
.vortex_expect("scalar_at should succeed in conformance test")
);
}
}
fn test_take_repeated_indices(array: &ArrayRef) {
let len = array.len();
if len == 0 {
return;
}
let indices = PrimitiveArray::from_iter([0u64, 0, 0]).into_array();
let taken = array
.take(indices)
.vortex_expect("take should succeed in conformance test");
assert_eq!(taken.len(), 3);
for i in 0..3 {
assert_eq!(
taken
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test"),
array
.scalar_at(0)
.vortex_expect("scalar_at should succeed in conformance test")
);
}
}
fn test_mask_filter_null_consistency(array: &ArrayRef) {
let len = array.len();
if len < 3 {
return;
}
let mask_pattern: Vec<bool> = (0..len).map(|i| i % 2 == 0).collect();
let mask_array = Mask::from_iter(mask_pattern);
let masked = array
.clone()
.mask((!&mask_array).into_array())
.vortex_expect("mask should succeed in conformance test");
let filter_pattern: Vec<bool> = (0..len).map(|i| i % 2 != 0).collect();
let filter_mask = Mask::from_iter(filter_pattern);
let filtered = masked
.filter(filter_mask.clone())
.vortex_expect("filter should succeed in conformance test");
let direct_filtered = array
.filter(filter_mask)
.vortex_expect("filter should succeed in conformance test");
assert_eq!(filtered.len(), direct_filtered.len());
for i in 0..filtered.len() {
assert_eq!(
filtered
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test"),
direct_filtered
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test")
);
}
}
fn test_empty_operations_consistency(array: &ArrayRef) {
let len = array.len();
let empty_filter = array
.filter(Mask::new_false(len))
.vortex_expect("filter should succeed in conformance test");
assert_eq!(empty_filter.len(), 0);
assert_eq!(empty_filter.dtype(), array.dtype());
let empty_indices = PrimitiveArray::empty::<u64>(Nullability::NonNullable).into_array();
let empty_take = array
.take(empty_indices)
.vortex_expect("take should succeed in conformance test");
assert_eq!(empty_take.len(), 0);
assert_eq!(empty_take.dtype(), array.dtype());
if len > 0 {
let empty_slice = array
.slice(0..0)
.vortex_expect("slice should succeed in conformance test");
assert_eq!(empty_slice.len(), 0);
assert_eq!(empty_slice.dtype(), array.dtype());
}
}
fn test_take_preserves_properties(array: &ArrayRef) {
let len = array.len();
if len == 0 {
return;
}
let indices = PrimitiveArray::from_iter((0..len).map(|i| i as u64)).into_array();
let taken = array
.take(indices)
.vortex_expect("take should succeed in conformance test");
assert_eq!(taken.len(), array.len());
assert_eq!(taken.dtype(), array.dtype());
for i in 0..len {
assert_eq!(
taken
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test"),
array
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test")
);
}
}
fn test_nullable_indices_consistency(array: &ArrayRef) {
let len = array.len();
if len < 3 {
return; }
let indices = PrimitiveArray::from_option_iter([Some(0u64), None, Some(2u64)]).into_array();
let taken = array
.take(indices)
.vortex_expect("take should succeed in conformance test");
assert_eq!(
taken.len(),
3,
"Take with nullable indices should produce array of length 3, got {}",
taken.len()
);
assert!(
taken.dtype().is_nullable(),
"Take with nullable indices should produce nullable array, but dtype is {:?}",
taken.dtype()
);
let expected_0 = array
.scalar_at(0)
.vortex_expect("scalar_at should succeed in conformance test")
.into_nullable();
let actual_0 = taken
.scalar_at(0)
.vortex_expect("scalar_at should succeed in conformance test");
assert_eq!(
actual_0, expected_0,
"Take with nullable indices: element at position 0 should be from array index 0. \
Expected: {expected_0:?}, Actual: {actual_0:?}"
);
let actual_1 = taken
.scalar_at(1)
.vortex_expect("scalar_at should succeed in conformance test");
assert!(
actual_1.is_null(),
"Take with nullable indices: element at position 1 should be null, but got {actual_1:?}"
);
let expected_2 = array
.scalar_at(2)
.vortex_expect("scalar_at should succeed in conformance test")
.into_nullable();
let actual_2 = taken
.scalar_at(2)
.vortex_expect("scalar_at should succeed in conformance test");
assert_eq!(
actual_2, expected_2,
"Take with nullable indices: element at position 2 should be from array index 2. \
Expected: {expected_2:?}, Actual: {actual_2:?}"
);
}
fn test_large_array_consistency(array: &ArrayRef) {
let len = array.len();
if len < 1000 {
return;
}
let indices: Vec<u64> = (0..len).step_by(10).map(|i| i as u64).collect();
let indices_array = PrimitiveArray::from_iter(indices).into_array();
let taken = array
.take(indices_array)
.vortex_expect("take should succeed in conformance test");
let mask_pattern: Vec<bool> = (0..len).map(|i| i % 10 == 0).collect();
let mask = Mask::from_iter(mask_pattern);
let filtered = array
.filter(mask)
.vortex_expect("filter should succeed in conformance test");
assert_eq!(taken.len(), filtered.len());
for i in 0..taken.len() {
assert_eq!(
taken
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test"),
filtered
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test")
);
}
}
fn test_comparison_inverse_consistency(array: &ArrayRef) {
let len = array.len();
if len == 0 {
return;
}
match array.dtype() {
DType::Null | DType::Extension(_) | DType::Struct(..) | DType::List(..) => return,
_ => {}
}
let test_scalar = if len == 0 {
return;
} else {
array
.scalar_at(len / 2)
.vortex_expect("scalar_at should succeed in conformance test")
};
let const_array = crate::arrays::ConstantArray::new(test_scalar, len);
if let (Ok(eq_result), Ok(neq_result)) = (
array
.clone()
.binary(const_array.clone().into_array(), Operator::Eq),
array
.clone()
.binary(const_array.clone().into_array(), Operator::NotEq),
) {
let inverted_eq = eq_result
.not()
.vortex_expect("not should succeed in conformance test");
assert_eq!(
inverted_eq.len(),
neq_result.len(),
"Inverted Eq should have same length as NotEq"
);
for i in 0..inverted_eq.len() {
let inv_val = inverted_eq
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test");
let neq_val = neq_result
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test");
assert_eq!(
inv_val, neq_val,
"At index {i}: NOT(Eq) should equal NotEq. \
NOT(Eq) = {inv_val:?}, NotEq = {neq_val:?}"
);
}
}
if let (Ok(gt_result), Ok(lte_result)) = (
array
.clone()
.binary(const_array.clone().into_array(), Operator::Gt),
array
.clone()
.binary(const_array.clone().into_array(), Operator::Lte),
) {
let inverted_gt = gt_result
.not()
.vortex_expect("not should succeed in conformance test");
for i in 0..inverted_gt.len() {
let inv_val = inverted_gt
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test");
let lte_val = lte_result
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test");
assert_eq!(
inv_val, lte_val,
"At index {i}: NOT(Gt) should equal Lte. \
NOT(Gt) = {inv_val:?}, Lte = {lte_val:?}"
);
}
}
if let (Ok(lt_result), Ok(gte_result)) = (
array
.clone()
.binary(const_array.clone().into_array(), Operator::Lt),
array
.clone()
.binary(const_array.into_array(), Operator::Gte),
) {
let inverted_lt = lt_result
.not()
.vortex_expect("not should succeed in conformance test");
for i in 0..inverted_lt.len() {
let inv_val = inverted_lt
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test");
let gte_val = gte_result
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test");
assert_eq!(
inv_val, gte_val,
"At index {i}: NOT(Lt) should equal Gte. \
NOT(Lt) = {inv_val:?}, Gte = {gte_val:?}"
);
}
}
}
fn test_comparison_symmetry_consistency(array: &ArrayRef) {
let len = array.len();
if len == 0 {
return;
}
match array.dtype() {
DType::Null | DType::Extension(_) | DType::Struct(..) | DType::List(..) => return,
_ => {}
}
let test_scalar = if len == 2 {
return;
} else {
array
.scalar_at(len / 2)
.vortex_expect("scalar_at should succeed in conformance test")
};
let const_array = crate::arrays::ConstantArray::new(test_scalar, len);
if let (Ok(arr_gt_scalar), Ok(scalar_lt_arr)) = (
array
.clone()
.binary(const_array.clone().into_array(), Operator::Gt),
const_array
.clone()
.into_array()
.binary(array.clone(), Operator::Lt),
) {
assert_eq!(
arr_gt_scalar.len(),
scalar_lt_arr.len(),
"Symmetric comparisons should have same length"
);
for i in 0..arr_gt_scalar.len() {
let arr_gt = arr_gt_scalar
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test");
let scalar_lt = scalar_lt_arr
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test");
assert_eq!(
arr_gt, scalar_lt,
"At index {i}: (array > scalar) should equal (scalar < array). \
array > scalar = {arr_gt:?}, scalar < array = {scalar_lt:?}"
);
}
}
if let (Ok(arr_eq_scalar), Ok(scalar_eq_arr)) = (
array
.clone()
.binary(const_array.clone().into_array(), Operator::Eq),
const_array.into_array().binary(array.clone(), Operator::Eq),
) {
for i in 0..arr_eq_scalar.len() {
let arr_eq = arr_eq_scalar
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test");
let scalar_eq = scalar_eq_arr
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test");
assert_eq!(
arr_eq, scalar_eq,
"At index {i}: (array == scalar) should equal (scalar == array). \
array == scalar = {arr_eq:?}, scalar == array = {scalar_eq:?}"
);
}
}
}
fn test_boolean_demorgan_consistency(array: &ArrayRef) {
if !matches!(array.dtype(), DType::Bool(_)) {
return;
}
let bool_mask = {
let mask_pattern: Vec<bool> = (0..array.len()).map(|i| i % 3 == 0).collect();
BoolArray::from_iter(mask_pattern)
};
let bool_mask = bool_mask.into_array();
if let (Ok(a_and_b), Ok(not_a), Ok(not_b)) = (
array.clone().binary(bool_mask.clone(), Operator::And),
array.not(),
bool_mask.not(),
) {
let not_a_and_b = a_and_b
.not()
.vortex_expect("not should succeed in conformance test");
let not_a_or_not_b = not_a
.binary(not_b, Operator::Or)
.vortex_expect("or should succeed in conformance test");
assert_eq!(
not_a_and_b.len(),
not_a_or_not_b.len(),
"De Morgan's law results should have same length"
);
for i in 0..not_a_and_b.len() {
let left = not_a_and_b
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test");
let right = not_a_or_not_b
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test");
assert_eq!(
left, right,
"De Morgan's first law failed at index {i}: \
NOT(A AND B) = {left:?}, (NOT A) OR (NOT B) = {right:?}"
);
}
}
if let (Ok(a_or_b), Ok(not_a), Ok(not_b)) = (
array.clone().binary(bool_mask.clone(), Operator::Or),
array.not(),
bool_mask.not(),
) {
let not_a_or_b = a_or_b
.not()
.vortex_expect("not should succeed in conformance test");
let not_a_and_not_b = not_a
.binary(not_b, Operator::And)
.vortex_expect("and should succeed in conformance test");
for i in 0..not_a_or_b.len() {
let left = not_a_or_b
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test");
let right = not_a_and_not_b
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test");
assert_eq!(
left, right,
"De Morgan's second law failed at index {i}: \
NOT(A OR B) = {left:?}, (NOT A) AND (NOT B) = {right:?}"
);
}
}
}
fn test_slice_aggregate_consistency(array: &ArrayRef) {
use crate::aggregate_fn::fns::min_max::min_max;
use crate::aggregate_fn::fns::nan_count::nan_count;
use crate::aggregate_fn::fns::sum::sum;
use crate::dtype::DType;
let mut ctx = LEGACY_SESSION.create_execution_ctx();
let len = array.len();
if len < 5 {
return; }
let start = 1;
let end = (len - 1).min(start + 10);
let sliced = array
.slice(start..end)
.vortex_expect("slice should succeed in conformance test");
let canonical = array.to_canonical().vortex_expect("to_canonical failed");
let canonical_sliced = canonical
.into_array()
.slice(start..end)
.vortex_expect("slice should succeed in conformance test");
let sliced_invalid_count = sliced
.invalid_count()
.vortex_expect("invalid_count should succeed in conformance test");
let canonical_invalid_count = canonical_sliced
.invalid_count()
.vortex_expect("invalid_count should succeed in conformance test");
assert_eq!(
sliced_invalid_count, canonical_invalid_count,
"null_count on sliced array should match canonical. \
Sliced: {sliced_invalid_count}, Canonical: {canonical_invalid_count}",
);
if !matches!(array.dtype(), DType::Primitive(..)) {
return;
}
if let (Ok(slice_sum), Ok(canonical_sum)) =
(sum(&sliced, &mut ctx), sum(&canonical_sliced, &mut ctx))
{
assert_eq!(
slice_sum, canonical_sum,
"sum on sliced array should match canonical. \
Sliced: {slice_sum:?}, Canonical: {canonical_sum:?}"
);
}
if let (Ok(slice_minmax), Ok(canonical_minmax)) = (
min_max(&sliced, &mut ctx),
min_max(&canonical_sliced, &mut ctx),
) {
match (slice_minmax, canonical_minmax) {
(Some(s_result), Some(c_result)) => {
assert_eq!(
s_result.min, c_result.min,
"min on sliced array should match canonical. \
Sliced: {:?}, Canonical: {:?}",
s_result.min, c_result.min
);
assert_eq!(
s_result.max, c_result.max,
"max on sliced array should match canonical. \
Sliced: {:?}, Canonical: {:?}",
s_result.max, c_result.max
);
}
(None, None) => {} _ => vortex_panic!("min_max results don't match"),
}
}
if array.dtype().is_float()
&& let (Ok(slice_nan_count), Ok(canonical_nan_count)) = (
nan_count(&sliced, &mut ctx),
nan_count(&canonical_sliced, &mut ctx),
)
{
assert_eq!(
slice_nan_count, canonical_nan_count,
"nan_count on sliced array should match canonical. \
Sliced: {slice_nan_count}, Canonical: {canonical_nan_count}"
);
}
}
fn test_cast_slice_consistency(array: &ArrayRef) {
let len = array.len();
if len < 5 {
return; }
let start = 2;
let end = 7.min(len - 2).max(start + 1);
let canonical = array.to_canonical().vortex_expect("to_canonical failed");
let target_dtypes = match array.dtype() {
DType::Null => vec![],
DType::Bool(nullability) => vec![
DType::Primitive(PType::U8, *nullability),
DType::Primitive(PType::I32, *nullability),
],
DType::Primitive(ptype, nullability) => {
let mut targets = vec![];
let opposite_nullability = match nullability {
Nullability::NonNullable => Nullability::Nullable,
Nullability::Nullable => Nullability::NonNullable,
};
targets.push(DType::Primitive(*ptype, opposite_nullability));
match ptype {
PType::U8 => {
targets.push(DType::Primitive(PType::U16, *nullability));
targets.push(DType::Primitive(PType::I16, *nullability));
}
PType::U16 => {
targets.push(DType::Primitive(PType::U32, *nullability));
targets.push(DType::Primitive(PType::I32, *nullability));
}
PType::U32 => {
targets.push(DType::Primitive(PType::U64, *nullability));
targets.push(DType::Primitive(PType::I64, *nullability));
}
PType::U64 => {
targets.push(DType::Primitive(PType::F64, *nullability));
}
PType::I8 => {
targets.push(DType::Primitive(PType::I16, *nullability));
targets.push(DType::Primitive(PType::F32, *nullability));
}
PType::I16 => {
targets.push(DType::Primitive(PType::I32, *nullability));
targets.push(DType::Primitive(PType::F32, *nullability));
}
PType::I32 => {
targets.push(DType::Primitive(PType::I64, *nullability));
targets.push(DType::Primitive(PType::F64, *nullability));
}
PType::I64 => {
targets.push(DType::Primitive(PType::F64, *nullability));
}
PType::F16 => {
targets.push(DType::Primitive(PType::F32, *nullability));
}
PType::F32 => {
targets.push(DType::Primitive(PType::F64, *nullability));
targets.push(DType::Primitive(PType::I32, *nullability));
}
PType::F64 => {
targets.push(DType::Primitive(PType::I64, *nullability));
}
}
targets
}
DType::Utf8(nullability) => {
let opposite = match nullability {
Nullability::NonNullable => Nullability::Nullable,
Nullability::Nullable => Nullability::NonNullable,
};
vec![DType::Utf8(opposite), DType::Binary(*nullability)]
}
DType::Binary(nullability) => {
let opposite = match nullability {
Nullability::NonNullable => Nullability::Nullable,
Nullability::Nullable => Nullability::NonNullable,
};
vec![
DType::Binary(opposite),
DType::Utf8(*nullability), ]
}
DType::Decimal(decimal_type, nullability) => {
let opposite = match nullability {
Nullability::NonNullable => Nullability::Nullable,
Nullability::Nullable => Nullability::NonNullable,
};
vec![DType::Decimal(*decimal_type, opposite)]
}
DType::Struct(fields, nullability) => {
let opposite = match nullability {
Nullability::NonNullable => Nullability::Nullable,
Nullability::Nullable => Nullability::NonNullable,
};
vec![DType::Struct(fields.clone(), opposite)]
}
DType::List(element_type, nullability) => {
let opposite = match nullability {
Nullability::NonNullable => Nullability::Nullable,
Nullability::Nullable => Nullability::NonNullable,
};
vec![DType::List(Arc::clone(element_type), opposite)]
}
DType::FixedSizeList(element_type, list_size, nullability) => {
let opposite = match nullability {
Nullability::NonNullable => Nullability::Nullable,
Nullability::Nullable => Nullability::NonNullable,
};
vec![DType::FixedSizeList(
Arc::clone(element_type),
*list_size,
opposite,
)]
}
DType::Extension(_) => vec![], DType::Variant(_) => unimplemented!(),
};
for target_dtype in target_dtypes {
let sliced = array
.slice(start..end)
.vortex_expect("slice should succeed in conformance test");
let slice_then_cast = match sliced
.cast(target_dtype.clone())
.and_then(|a| a.to_canonical().map(|c| c.into_array()))
{
Ok(result) => result,
Err(_) => continue, };
assert_eq!(
slice_then_cast.len(),
end - start,
"Sliced and casted array should have length {}, but has {}",
end - start,
slice_then_cast.len()
);
for i in 0..slice_then_cast.len() {
let slice_cast_val = slice_then_cast
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test");
let canonical_val = canonical
.clone()
.into_array()
.scalar_at(start + i)
.vortex_expect("scalar_at should succeed in conformance test");
let expected_val = match canonical_val.cast(&target_dtype) {
Ok(val) => val,
Err(_) => {
break;
}
};
assert_eq!(
slice_cast_val,
expected_val,
"Cast of sliced array produced incorrect value at index {i}. \
Got: {slice_cast_val:?}, Expected: {expected_val:?} \
(canonical value at index {}: {canonical_val:?})\n\
This likely indicates the array encoding doesn't preserve offset information during cast.",
start + i
);
}
let casted = match array
.clone()
.cast(target_dtype.clone())
.and_then(|a| a.to_canonical().map(|c| c.into_array()))
{
Ok(result) => result,
Err(_) => continue, };
let cast_then_slice = casted
.slice(start..end)
.vortex_expect("slice should succeed in conformance test");
assert_eq!(
slice_then_cast.len(),
cast_then_slice.len(),
"Slice-then-cast and cast-then-slice should produce arrays of the same length"
);
for i in 0..slice_then_cast.len() {
let slice_cast_val = slice_then_cast
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test");
let cast_slice_val = cast_then_slice
.scalar_at(i)
.vortex_expect("scalar_at should succeed in conformance test");
assert_eq!(
slice_cast_val, cast_slice_val,
"Slice-then-cast and cast-then-slice produced different values at index {i}. \
Slice-then-cast: {slice_cast_val:?}, Cast-then-slice: {cast_slice_val:?}"
);
}
}
}
pub fn test_array_consistency(array: &ArrayRef) {
test_filter_take_consistency(array);
test_double_mask_consistency(array);
test_slice_filter_consistency(array);
test_take_slice_consistency(array);
test_cast_slice_consistency(array);
test_boolean_demorgan_consistency(array);
test_comparison_inverse_consistency(array);
test_comparison_symmetry_consistency(array);
test_slice_aggregate_consistency(array);
test_filter_identity(array);
test_mask_identity(array);
test_take_preserves_properties(array);
test_filter_preserves_order(array);
test_take_repeated_indices(array);
test_mask_filter_null_consistency(array);
test_nullable_indices_consistency(array);
test_empty_operations_consistency(array);
test_large_array_consistency(array);
}