Skip to content

Commit

Permalink
merge
Browse files Browse the repository at this point in the history
  • Loading branch information
gatesn committed Mar 4, 2024
1 parent bc2c5dc commit 1a7efaf
Show file tree
Hide file tree
Showing 12 changed files with 54 additions and 44 deletions.
6 changes: 3 additions & 3 deletions vortex-dict/src/compress.rs
Original file line number Diff line number Diff line change
Expand Up @@ -269,7 +269,7 @@ where
mod test {
use vortex::array::primitive::PrimitiveArray;
use vortex::array::varbin::VarBinArray;
use vortex::array::Array;
use vortex::compute::scalar_at::scalar_at;

use crate::compress::{dict_encode_typed_primitive, dict_encode_varbin};

Expand Down Expand Up @@ -298,8 +298,8 @@ mod test {
assert!(!codes.is_valid(2));
assert!(!codes.is_valid(5));
assert!(!codes.is_valid(7));
assert_eq!(values.scalar_at(0), Ok(1.into()));
assert_eq!(values.scalar_at(2), Ok(3.into()));
assert_eq!(scalar_at(values.as_ref(), 0), Ok(1.into()));
assert_eq!(scalar_at(values.as_ref(), 2), Ok(3.into()));
}

#[test]
Expand Down
9 changes: 5 additions & 4 deletions vortex-ree/src/ree.rs
Original file line number Diff line number Diff line change
Expand Up @@ -315,6 +315,7 @@ mod test {
use arrow::array::types::Int32Type;
use itertools::Itertools;
use vortex::array::Array;
use vortex::compute::scalar_at::scalar_at;

use crate::REEArray;
use vortex::dtype::{DType, IntWidth, Nullability, Signedness};
Expand All @@ -331,10 +332,10 @@ mod test {
// 0, 1 => 1
// 2, 3, 4 => 2
// 5, 6, 7, 8, 9 => 3
assert_eq!(arr.scalar_at(0).unwrap().try_into(), Ok(1));
assert_eq!(arr.scalar_at(2).unwrap().try_into(), Ok(2));
assert_eq!(arr.scalar_at(5).unwrap().try_into(), Ok(3));
assert_eq!(arr.scalar_at(9).unwrap().try_into(), Ok(3));
assert_eq!(scalar_at(arr.as_ref(), 0).unwrap().try_into(), Ok(1));
assert_eq!(scalar_at(arr.as_ref(), 2).unwrap().try_into(), Ok(2));
assert_eq!(scalar_at(arr.as_ref(), 5).unwrap().try_into(), Ok(3));
assert_eq!(scalar_at(arr.as_ref(), 9).unwrap().try_into(), Ok(3));
}

#[test]
Expand Down
11 changes: 6 additions & 5 deletions vortex-roaring/src/boolean/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -155,6 +155,7 @@ impl Encoding for RoaringBoolEncoding {
mod test {
use vortex::array::bool::BoolArray;
use vortex::array::Array;
use vortex::compute::scalar_at::scalar_at;
use vortex::error::VortexResult;
use vortex::scalar::Scalar;

Expand All @@ -172,17 +173,17 @@ mod test {
}

#[test]
pub fn scalar_at() -> VortexResult<()> {
pub fn test_scalar_at() -> VortexResult<()> {
let bool: &dyn Array = &BoolArray::from(vec![true, false, true, true]);
let array = RoaringBoolArray::encode(bool)?;

let truthy: Box<dyn Scalar> = true.into();
let falsy: Box<dyn Scalar> = false.into();

assert_eq!(array.scalar_at(0)?, truthy);
assert_eq!(array.scalar_at(1)?, falsy);
assert_eq!(array.scalar_at(2)?, truthy);
assert_eq!(array.scalar_at(3)?, truthy);
assert_eq!(scalar_at(array.as_ref(), 0)?, truthy);
assert_eq!(scalar_at(array.as_ref(), 1)?, falsy);
assert_eq!(scalar_at(array.as_ref(), 2)?, truthy);
assert_eq!(scalar_at(array.as_ref(), 3)?, truthy);

Ok(())
}
Expand Down
7 changes: 4 additions & 3 deletions vortex-roaring/src/integer/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -157,17 +157,18 @@ impl Encoding for RoaringIntEncoding {
mod test {
use vortex::array::primitive::PrimitiveArray;
use vortex::array::Array;
use vortex::compute::scalar_at::scalar_at;
use vortex::error::VortexResult;

use crate::RoaringIntArray;

#[test]
pub fn scalar_at() -> VortexResult<()> {
pub fn test_scalar_at() -> VortexResult<()> {
let ints: &dyn Array = &PrimitiveArray::from_vec::<u32>(vec![2, 12, 22, 32]);
let array = RoaringIntArray::encode(ints)?;

assert_eq!(array.scalar_at(0), Ok(2u32.into()));
assert_eq!(array.scalar_at(1), Ok(12u32.into()));
assert_eq!(scalar_at(array.as_ref(), 0), Ok(2u32.into()));
assert_eq!(scalar_at(array.as_ref(), 1), Ok(12u32.into()));

Ok(())
}
Expand Down
6 changes: 3 additions & 3 deletions vortex/src/array/bool/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -234,9 +234,9 @@ mod test {
.slice(1, 4)
.unwrap();
assert_eq!(arr.len(), 3);
assert_eq!(arr.scalar_at(0).unwrap().try_into(), Ok(true));
assert_eq!(arr.scalar_at(1).unwrap().try_into(), Ok(false));
assert_eq!(arr.scalar_at(2).unwrap().try_into(), Ok(false));
assert_eq!(scalar_at(arr.as_ref(), 0).unwrap().try_into(), Ok(true));
assert_eq!(scalar_at(arr.as_ref(), 1).unwrap().try_into(), Ok(false));
assert_eq!(scalar_at(arr.as_ref(), 2).unwrap().try_into(), Ok(false));
}

#[test]
Expand Down
3 changes: 2 additions & 1 deletion vortex/src/array/primitive/compress.rs
Original file line number Diff line number Diff line change
Expand Up @@ -32,12 +32,13 @@ mod test {
use crate::array::primitive::PrimitiveArray;
use crate::array::Encoding;
use crate::compress::CompressCtx;
use crate::compute::scalar_at::scalar_at;

#[test]
pub fn compress_constant() {
let arr = PrimitiveArray::from_vec(vec![1, 1, 1, 1]);
let res = CompressCtx::default().compress(arr.as_ref(), None);
assert_eq!(res.encoding().id(), ConstantEncoding.id());
assert_eq!(res.scalar_at(3).unwrap().try_into(), Ok(1));
assert_eq!(scalar_at(res.as_ref(), 3).unwrap().try_into(), Ok(1));
}
}
12 changes: 6 additions & 6 deletions vortex/src/array/primitive/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -319,9 +319,9 @@ mod test {
);

// Ensure we can fetch the scalar at the given index.
assert_eq!(arr.scalar_at(0).unwrap().try_into(), Ok(1));
assert_eq!(arr.scalar_at(1).unwrap().try_into(), Ok(2));
assert_eq!(arr.scalar_at(2).unwrap().try_into(), Ok(3));
assert_eq!(scalar_at(arr.as_ref(), 0).unwrap().try_into(), Ok(1));
assert_eq!(scalar_at(arr.as_ref(), 1).unwrap().try_into(), Ok(2));
assert_eq!(scalar_at(arr.as_ref(), 2).unwrap().try_into(), Ok(3));
}

#[test]
Expand All @@ -330,8 +330,8 @@ mod test {
.slice(1, 4)
.unwrap();
assert_eq!(arr.len(), 3);
assert_eq!(arr.scalar_at(0).unwrap().try_into(), Ok(2));
assert_eq!(arr.scalar_at(1).unwrap().try_into(), Ok(3));
assert_eq!(arr.scalar_at(2).unwrap().try_into(), Ok(4));
assert_eq!(scalar_at(arr.as_ref(), 0).unwrap().try_into(), Ok(2));
assert_eq!(scalar_at(arr.as_ref(), 1).unwrap().try_into(), Ok(3));
assert_eq!(scalar_at(arr.as_ref(), 2).unwrap().try_into(), Ok(4));
}
}
22 changes: 13 additions & 9 deletions vortex/src/array/sparse/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -231,6 +231,7 @@ mod test {

use crate::array::sparse::SparseArray;
use crate::array::Array;
use crate::compute::scalar_at::scalar_at;
use crate::error::VortexError;

fn sparse_array() -> SparseArray {
Expand Down Expand Up @@ -291,23 +292,26 @@ mod test {
}

#[test]
pub fn scalar_at() {
pub fn test_scalar_at() {
assert_eq!(
usize::try_from(sparse_array().scalar_at(2).unwrap()).unwrap(),
usize::try_from(scalar_at(sparse_array().as_ref(), 2).unwrap()).unwrap(),
100
);
assert_eq!(
sparse_array().scalar_at(10).err().unwrap(),
scalar_at(sparse_array().as_ref(), 10).err().unwrap(),
VortexError::OutOfBounds(10, 0, 10)
);
}

#[test]
pub fn scalar_at_sliced() {
let sliced = sparse_array().slice(2, 7).unwrap();
assert_eq!(usize::try_from(sliced.scalar_at(0).unwrap()).unwrap(), 100);
assert_eq!(
sliced.scalar_at(5).err().unwrap(),
usize::try_from(scalar_at(sliced.as_ref(), 0).unwrap()).unwrap(),
100
);
assert_eq!(
scalar_at(sliced.as_ref(), 5).err().unwrap(),
VortexError::OutOfBounds(5, 0, 5)
);
}
Expand All @@ -316,21 +320,21 @@ mod test {
pub fn scalar_at_sliced_twice() {
let sliced_once = sparse_array().slice(1, 8).unwrap();
assert_eq!(
usize::try_from(sliced_once.scalar_at(1).unwrap()).unwrap(),
usize::try_from(scalar_at(sliced_once.as_ref(), 1).unwrap()).unwrap(),
100
);
assert_eq!(
sliced_once.scalar_at(7).err().unwrap(),
scalar_at(sliced_once.as_ref(), 7).err().unwrap(),
VortexError::OutOfBounds(7, 0, 7)
);

let sliced_twice = sliced_once.slice(1, 6).unwrap();
assert_eq!(
usize::try_from(sliced_twice.scalar_at(3).unwrap()).unwrap(),
usize::try_from(scalar_at(sliced_twice.as_ref(), 3).unwrap()).unwrap(),
200
);
assert_eq!(
sliced_twice.scalar_at(5).err().unwrap(),
scalar_at(sliced_twice.as_ref(), 5).err().unwrap(),
VortexError::OutOfBounds(5, 0, 5)
);
}
Expand Down
5 changes: 3 additions & 2 deletions vortex/src/array/typed/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -161,6 +161,7 @@ mod test {

use crate::array::typed::TypedArray;
use crate::array::Array;
use crate::compute::scalar_at::scalar_at;
use crate::dtype::{DType, Nullability, TimeUnit};
use crate::scalar::{LocalTimeScalar, PScalar, Scalar};

Expand All @@ -171,11 +172,11 @@ mod test {
DType::LocalTime(TimeUnit::Us, Nullability::NonNullable),
);
assert_eq!(
arr.scalar_at(0).unwrap().as_ref(),
scalar_at(arr.as_ref(), 0).unwrap().as_ref(),
&LocalTimeScalar::new(PScalar::U64(64_799_000_000), TimeUnit::Us) as &dyn Scalar
);
assert_eq!(
arr.scalar_at(1).unwrap().as_ref(),
scalar_at(arr.as_ref(), 1).unwrap().as_ref(),
&LocalTimeScalar::new(PScalar::U64(43_000_000_000), TimeUnit::Us) as &dyn Scalar
);
}
Expand Down
9 changes: 5 additions & 4 deletions vortex/src/array/varbin/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -383,6 +383,7 @@ mod test {
use crate::array::primitive::PrimitiveArray;
use crate::array::varbin::VarBinArray;
use crate::arrow::CombineChunks;
use crate::compute::scalar_at::scalar_at;
use crate::dtype::{DType, Nullability};

fn binary_array() -> VarBinArray {
Expand All @@ -402,12 +403,12 @@ mod test {
}

#[test]
pub fn scalar_at() {
pub fn test_scalar_at() {
let binary_arr = binary_array();
assert_eq!(binary_arr.len(), 2);
assert_eq!(binary_arr.scalar_at(0), Ok("hello world".into()));
assert_eq!(scalar_at(binary_arr.as_ref(), 0), Ok("hello world".into()));
assert_eq!(
binary_arr.scalar_at(1),
scalar_at(binary_arr.as_ref(), 1),
Ok("hello world this is a long string".into())
)
}
Expand All @@ -416,7 +417,7 @@ mod test {
pub fn slice() {
let binary_arr = binary_array().slice(1, 2).unwrap();
assert_eq!(
binary_arr.scalar_at(0),
scalar_at(binary_arr.as_ref(), 0),
Ok("hello world this is a long string".into())
);
}
Expand Down
6 changes: 3 additions & 3 deletions vortex/src/array/varbinview/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -392,9 +392,9 @@ mod test {
pub fn varbin_view() {
let binary_arr = binary_array();
assert_eq!(binary_arr.len(), 2);
assert_eq!(binary_arr.scalar_at(0), Ok("hello world".into()));
assert_eq!(scalar_at(binary_arr.as_ref(), 0), Ok("hello world".into()));
assert_eq!(
binary_arr.scalar_at(1),
scalar_at(binary_arr.as_ref(), 1),
Ok("hello world this is a long string".into())
)
}
Expand All @@ -403,7 +403,7 @@ mod test {
pub fn slice() {
let binary_arr = binary_array().slice(1, 2).unwrap();
assert_eq!(
binary_arr.scalar_at(0),
scalar_at(binary_arr.as_ref(), 0),
Ok("hello world this is a long string".into())
);
}
Expand Down
2 changes: 1 addition & 1 deletion vortex/src/compute/add.rs
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,6 @@ mod test {
let rhs = ConstantArray::new(47.into(), 100);
let result = add(&lhs, &rhs).unwrap();
assert_eq!(result.len(), 100);
// assert_eq!(result.scalar_at(0), 94);
// assert_eq!(scalar_at(result, 0), 94);
}
}

0 comments on commit 1a7efaf

Please sign in to comment.