Skip to content

Commit

Permalink
Added list to the fuzzer (#1712)
Browse files Browse the repository at this point in the history
  • Loading branch information
joseph-isaacs authored Dec 20, 2024
1 parent 9f3d8e1 commit 0cc9df2
Show file tree
Hide file tree
Showing 3 changed files with 50 additions and 7 deletions.
2 changes: 1 addition & 1 deletion vortex-array/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ vortex-datetime-dtype = { workspace = true }
vortex-dtype = { workspace = true, features = ["flatbuffers", "serde"] }
vortex-error = { workspace = true, features = ["flatbuffers", "flexbuffers"] }
vortex-flatbuffers = { workspace = true, features = ["array"] }
vortex-scalar = { workspace = true, features = ["flatbuffers", "serde"] }
vortex-scalar = { workspace = true, features = ["flatbuffers", "serde", "arbitrary"] }

[features]
arbitrary = ["dep:arbitrary", "vortex-dtype/arbitrary"]
Expand Down
53 changes: 48 additions & 5 deletions vortex-array/src/array/arbitrary.rs
Original file line number Diff line number Diff line change
@@ -1,14 +1,20 @@
use std::iter;
use std::sync::Arc;

use arbitrary::{Arbitrary, Result, Unstructured};
use arrow_buffer::BooleanBuffer;
use builders::ListBuilder;
use num_traits::{AsPrimitive, PrimInt};
use vortex_dtype::{DType, NativePType, Nullability, PType};
use vortex_error::{VortexExpect, VortexUnwrap};
use vortex_scalar::arbitrary::random_scalar;
use vortex_scalar::Scalar;

use super::{BoolArray, ChunkedArray, NullArray, PrimitiveArray, StructArray};
use crate::array::{VarBinArray, VarBinViewArray};
use crate::builders::ArrayBuilder;
use crate::validity::Validity;
use crate::{ArrayDType, ArrayData, IntoArrayData as _, IntoArrayVariant};
use crate::{builders, ArrayDType, ArrayData, IntoArrayData as _, IntoArrayVariant};

impl<'a> Arbitrary<'a> for ArrayData {
fn arbitrary(u: &mut Unstructured<'a>) -> Result<Self> {
Expand Down Expand Up @@ -81,10 +87,7 @@ fn random_array(u: &mut Unstructured, dtype: &DType, len: Option<usize>) -> Resu
.vortex_unwrap()
.into_array())
}
// TOOD(joe): add arbitrary list
DType::List(..) => {
todo!("List arrays are not implemented")
}
DType::List(ldt, n) => random_list(u, ldt, n),
DType::Extension(..) => {
todo!("Extension arrays are not implemented")
}
Expand All @@ -102,6 +105,46 @@ fn random_array(u: &mut Unstructured, dtype: &DType, len: Option<usize>) -> Resu
}
}

fn random_list(u: &mut Unstructured, ldt: &Arc<DType>, n: &Nullability) -> Result<ArrayData> {
match u.int_in_range(0..=5)? {
0 => random_list_offset::<i16>(u, ldt, n),
1 => random_list_offset::<i32>(u, ldt, n),
2 => random_list_offset::<i64>(u, ldt, n),
3 => random_list_offset::<u16>(u, ldt, n),
4 => random_list_offset::<u32>(u, ldt, n),
5 => random_list_offset::<u64>(u, ldt, n),
_ => unreachable!("int_in_range returns a value in the above range"),
}
}

fn random_list_offset<O>(
u: &mut Unstructured,
ldt: &Arc<DType>,
n: &Nullability,
) -> Result<ArrayData>
where
O: PrimInt + NativePType,
Scalar: From<O>,
usize: AsPrimitive<O>,
{
let list_len = u.int_in_range(0..=20)?;
let mut builder = ListBuilder::<O>::with_capacity(ldt.clone(), *n, 1);
for _ in 0..list_len {
if matches!(n, Nullability::Nullable) || u.arbitrary::<bool>()? {
let elem_len = u.int_in_range(0..=20)?;
let elem = (0..elem_len)
.map(|_| random_scalar(u, ldt))
.collect::<Result<Vec<_>>>()?;
builder
.append_value(Scalar::list(ldt.clone(), elem, *n).as_list())
.vortex_expect("can append value");
} else {
builder.append_null();
}
}
Ok(builder.finish().vortex_expect("builder cannot error"))
}

fn split_number_into_parts(n: usize, parts: usize) -> Vec<usize> {
let reminder = n % parts;
let division = (n - reminder) / parts;
Expand Down
2 changes: 1 addition & 1 deletion vortex-array/src/builders/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ use std::any::Any;
pub use binary::*;
pub use bool::*;
pub use extension::*;
pub use list::*;
pub use null::*;
pub use primitive::*;
pub use utf8::*;
Expand All @@ -22,7 +23,6 @@ use vortex_scalar::{
Utf8Scalar,
};

use crate::builders::list::ListBuilder;
use crate::builders::struct_::StructBuilder;
use crate::ArrayData;

Expand Down

0 comments on commit 0cc9df2

Please sign in to comment.