diff --git a/encodings/alp/src/alp_rd/compute/take.rs b/encodings/alp/src/alp_rd/compute/take.rs index b2a395c0813..7b123f19eca 100644 --- a/encodings/alp/src/alp_rd/compute/take.rs +++ b/encodings/alp/src/alp_rd/compute/take.rs @@ -62,7 +62,6 @@ mod test { use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::compute::conformance::take::test_take_conformance; - use vortex_array::compute::take; use crate::ALPRDFloat; use crate::RDEncoder; @@ -86,7 +85,8 @@ mod test { .is_unsigned_int() ); - let taken = take(encoded.as_ref(), buffer![0, 2].into_array().as_ref()) + let taken = encoded + .take(buffer![0, 2].into_array()) .unwrap() .to_primitive(); @@ -109,12 +109,10 @@ mod test { .is_unsigned_int() ); - let taken = take( - encoded.as_ref(), - PrimitiveArray::from_option_iter([Some(0), Some(2), None]).as_ref(), - ) - .unwrap() - .to_primitive(); + let taken = encoded + .take(PrimitiveArray::from_option_iter([Some(0), Some(2), None]).to_array()) + .unwrap() + .to_primitive(); assert_arrays_eq!( taken, diff --git a/encodings/fastlanes/benches/bitpacking_take.rs b/encodings/fastlanes/benches/bitpacking_take.rs index 71a72894c4f..3c688598721 100644 --- a/encodings/fastlanes/benches/bitpacking_take.rs +++ b/encodings/fastlanes/benches/bitpacking_take.rs @@ -9,9 +9,9 @@ use rand::Rng; use rand::SeedableRng; use rand::distr::Uniform; use rand::prelude::StdRng; +use vortex_array::Array; use vortex_array::IntoArray as _; use vortex_array::arrays::PrimitiveArray; -use vortex_array::compute::take; use vortex_array::compute::warm_up_vtables; use vortex_array::validity::Validity; use vortex_buffer::Buffer; @@ -32,7 +32,7 @@ fn take_10_stratified(bencher: Bencher) { bencher .with_inputs(|| (&packed, &indices)) - .bench_refs(|(packed, indices)| take(packed.as_ref(), indices.as_ref()).unwrap()) + .bench_refs(|(packed, indices)| packed.take(indices.to_array()).unwrap()) } #[divan::bench] @@ -44,7 +44,7 @@ fn take_10_contiguous(bencher: Bencher) { bencher .with_inputs(|| (&packed, &indices)) - .bench_refs(|(packed, indices)| take(packed.as_ref(), indices.as_ref()).unwrap()) + .bench_refs(|(packed, indices)| packed.take(indices.to_array()).unwrap()) } #[divan::bench] @@ -59,7 +59,7 @@ fn take_10k_random(bencher: Bencher) { bencher .with_inputs(|| (&packed, &indices)) - .bench_refs(|(packed, indices)| take(packed.as_ref(), indices.as_ref()).unwrap()) + .bench_refs(|(packed, indices)| packed.take(indices.to_array()).unwrap()) } #[divan::bench] @@ -71,7 +71,7 @@ fn take_10k_contiguous(bencher: Bencher) { bencher .with_inputs(|| (&packed, &indices)) - .bench_refs(|(packed, indices)| take(packed.as_ref(), indices.as_ref()).unwrap()) + .bench_refs(|(packed, indices)| packed.take(indices.to_array()).unwrap()) } #[divan::bench] @@ -83,7 +83,7 @@ fn take_200k_dispersed(bencher: Bencher) { bencher .with_inputs(|| (&packed, &indices)) - .bench_refs(|(packed, indices)| take(packed.as_ref(), indices.as_ref()).unwrap()) + .bench_refs(|(packed, indices)| packed.take(indices.to_array()).unwrap()) } #[divan::bench] @@ -95,7 +95,7 @@ fn take_200k_first_chunk_only(bencher: Bencher) { bencher .with_inputs(|| (&packed, &indices)) - .bench_refs(|(packed, indices)| take(packed.as_ref(), indices.as_ref()).unwrap()) + .bench_refs(|(packed, indices)| packed.take(indices.to_array()).unwrap()) } fn fixture(len: usize, bits: usize) -> Buffer { @@ -134,7 +134,7 @@ fn patched_take_10_stratified(bencher: Bencher) { bencher .with_inputs(|| (&packed, &indices)) - .bench_refs(|(packed, indices)| take(packed.as_ref(), indices.as_ref()).unwrap()) + .bench_refs(|(packed, indices)| packed.take(indices.to_array()).unwrap()) } #[divan::bench] @@ -153,7 +153,7 @@ fn patched_take_10_contiguous(bencher: Bencher) { bencher .with_inputs(|| (&packed, &indices)) - .bench_refs(|(packed, indices)| take(packed.as_ref(), indices.as_ref()).unwrap()) + .bench_refs(|(packed, indices)| packed.take(indices.to_array()).unwrap()) } #[divan::bench] @@ -168,7 +168,7 @@ fn patched_take_10k_random(bencher: Bencher) { bencher .with_inputs(|| (&packed, &indices)) - .bench_refs(|(packed, indices)| take(packed.as_ref(), indices.as_ref()).unwrap()) + .bench_refs(|(packed, indices)| packed.take(indices.to_array()).unwrap()) } #[divan::bench] @@ -180,7 +180,7 @@ fn patched_take_10k_contiguous_not_patches(bencher: Bencher) { bencher .with_inputs(|| (&packed, &indices)) - .bench_refs(|(packed, indices)| take(packed.as_ref(), indices.as_ref()).unwrap()) + .bench_refs(|(packed, indices)| packed.take(indices.to_array()).unwrap()) } #[divan::bench] @@ -200,7 +200,7 @@ fn patched_take_10k_contiguous_patches(bencher: Bencher) { bencher .with_inputs(|| (&packed, &indices)) - .bench_refs(|(packed, indices)| take(packed.as_ref(), indices.as_ref()).unwrap()) + .bench_refs(|(packed, indices)| packed.take(indices.to_array()).unwrap()) } #[divan::bench] @@ -212,7 +212,7 @@ fn patched_take_200k_dispersed(bencher: Bencher) { bencher .with_inputs(|| (&packed, &indices)) - .bench_refs(|(packed, indices)| take(packed.as_ref(), indices.as_ref()).unwrap()) + .bench_refs(|(packed, indices)| packed.take(indices.to_array()).unwrap()) } #[divan::bench] @@ -224,7 +224,7 @@ fn patched_take_200k_first_chunk_only(bencher: Bencher) { bencher .with_inputs(|| (&packed, &indices)) - .bench_refs(|(packed, indices)| take(packed.as_ref(), indices.as_ref()).unwrap()) + .bench_refs(|(packed, indices)| packed.take(indices.to_array()).unwrap()) } #[divan::bench] @@ -243,5 +243,5 @@ fn patched_take_10k_adversarial(bencher: Bencher) { bencher .with_inputs(|| (&packed, &indices)) - .bench_refs(|(packed, indices)| take(packed.as_ref(), indices.as_ref()).unwrap()) + .bench_refs(|(packed, indices)| packed.take(indices.to_array()).unwrap()) } diff --git a/encodings/fastlanes/src/bitpacking/compute/take.rs b/encodings/fastlanes/src/bitpacking/compute/take.rs index 6a5b33495dc..2ce7c9583eb 100644 --- a/encodings/fastlanes/src/bitpacking/compute/take.rs +++ b/encodings/fastlanes/src/bitpacking/compute/take.rs @@ -155,7 +155,6 @@ mod test { use vortex_array::ToCanonical; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; - use vortex_array::compute::take; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::buffer; @@ -171,7 +170,7 @@ mod test { let unpacked = PrimitiveArray::from_iter((0..4096).map(|i| (i % 63) as u8)); let bitpacked = BitPackedArray::encode(unpacked.as_ref(), 6).unwrap(); - let primitive_result = take(bitpacked.as_ref(), &indices).unwrap(); + let primitive_result = bitpacked.take(indices.to_array()).unwrap(); assert_arrays_eq!( primitive_result, PrimitiveArray::from_iter([0u8, 62, 31, 33, 9, 18]) @@ -185,7 +184,7 @@ mod test { let indices = buffer![0, 2, 4, 6].into_array(); - let primitive_result = take(bitpacked.as_ref(), indices.as_ref()).unwrap(); + let primitive_result = bitpacked.take(indices.to_array()).unwrap(); assert_arrays_eq!(primitive_result, PrimitiveArray::from_iter([0u32, 2, 4, 6])); } @@ -198,7 +197,7 @@ mod test { let bitpacked = BitPackedArray::encode(unpacked.as_ref(), 6).unwrap(); let sliced = bitpacked.slice(128..2050).unwrap(); - let primitive_result = take(&sliced, &indices).unwrap(); + let primitive_result = sliced.take(indices.to_array()).unwrap(); assert_arrays_eq!(primitive_result, PrimitiveArray::from_iter([31u8, 33])); } @@ -215,7 +214,7 @@ mod test { let range = Uniform::new(0, values.len()).unwrap(); let random_indices = PrimitiveArray::from_iter(rng.sample_iter(range).take(10_000).map(|i| i as u32)); - let taken = take(packed.as_ref(), random_indices.as_ref()).unwrap(); + let taken = packed.take(random_indices.to_array()).unwrap(); // sanity check random_indices @@ -254,11 +253,9 @@ mod test { let start = BitPackedArray::encode(&buffer![1i32, 2i32, 3i32, 4i32].into_array(), 1).unwrap(); - let taken_primitive = take( - start.as_ref(), - PrimitiveArray::from_option_iter([Some(0u64), Some(1), None, Some(3)]).as_ref(), - ) - .unwrap(); + let taken_primitive = start + .take(PrimitiveArray::from_option_iter([Some(0u64), Some(1), None, Some(3)]).to_array()) + .unwrap(); assert_arrays_eq!( taken_primitive, PrimitiveArray::from_option_iter([Some(1i32), Some(2), None, Some(4)]) diff --git a/encodings/fastlanes/src/bitpacking/vtable/operations.rs b/encodings/fastlanes/src/bitpacking/vtable/operations.rs index 7d08800cbbc..95aed394760 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/operations.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/operations.rs @@ -36,7 +36,6 @@ mod test { use vortex_array::assert_arrays_eq; use vortex_array::assert_nth_scalar; use vortex_array::buffer::BufferHandle; - use vortex_array::compute::take; use vortex_array::patches::Patches; use vortex_array::session::ArraySession; use vortex_array::validity::Validity; @@ -175,7 +174,9 @@ mod test { // Chunk 2: physical indices 1024-2047, logical indices 102-1125 // Chunk 3: physical indices 2048-2060, logical indices 1126-1138 - let taken = take(&sliced, buffer![101i64, 1125, 1138].into_array().as_ref()).unwrap(); + let taken = sliced + .take(buffer![101i64, 1125, 1138].into_array()) + .unwrap(); assert_eq!(taken.len(), 3); } diff --git a/encodings/fsst/src/compute/mod.rs b/encodings/fsst/src/compute/mod.rs index b2657a73bae..53d835e3306 100644 --- a/encodings/fsst/src/compute/mod.rs +++ b/encodings/fsst/src/compute/mod.rs @@ -55,11 +55,11 @@ impl TakeExecute for FSSTVTable { #[cfg(test)] mod tests { use rstest::rstest; + use vortex_array::Array; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinArray; use vortex_array::compute::conformance::consistency::test_array_consistency; use vortex_array::compute::conformance::take::test_take_conformance; - use vortex_array::compute::take; use vortex_dtype::DType; use vortex_dtype::Nullability; @@ -76,14 +76,14 @@ mod tests { let idx1: PrimitiveArray = (0..1).collect(); assert_eq!( - take(fsst.as_ref(), idx1.as_ref()).unwrap().dtype(), + fsst.take(idx1.to_array()).unwrap().dtype(), &DType::Utf8(Nullability::NonNullable) ); let idx2: PrimitiveArray = PrimitiveArray::from_option_iter(vec![Some(0)]); assert_eq!( - take(fsst.as_ref(), idx2.as_ref()).unwrap().dtype(), + fsst.take(idx2.to_array()).unwrap().dtype(), &DType::Utf8(Nullability::Nullable) ); } diff --git a/encodings/runend/benches/run_end_compress.rs b/encodings/runend/benches/run_end_compress.rs index aeb0d9b9c73..24cabc6f155 100644 --- a/encodings/runend/benches/run_end_compress.rs +++ b/encodings/runend/benches/run_end_compress.rs @@ -8,7 +8,6 @@ use itertools::repeat_n; use vortex_array::Array; use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; -use vortex_array::compute::take; use vortex_array::compute::warm_up_vtables; use vortex_array::validity::Validity; use vortex_buffer::Buffer; @@ -97,5 +96,5 @@ fn take_indices(bencher: Bencher, (length, run_step): (usize, usize)) { bencher .with_inputs(|| (&source_array, &runend_array)) - .bench_refs(|(array, indices)| take(array.as_ref(), indices.as_ref()).unwrap()); + .bench_refs(|(array, indices)| array.take(indices.to_array()).unwrap()); } diff --git a/encodings/runend/src/compute/take.rs b/encodings/runend/src/compute/take.rs index e0a2dcdedf2..d8bfdd3fc1b 100644 --- a/encodings/runend/src/compute/take.rs +++ b/encodings/runend/src/compute/take.rs @@ -99,7 +99,6 @@ mod test { use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::compute::conformance::take::test_take_conformance; - use vortex_array::compute::take; use vortex_buffer::buffer; use crate::RunEndArray; @@ -110,18 +109,14 @@ mod test { #[test] fn ree_take() { - let taken = take( - ree_array().as_ref(), - buffer![9, 8, 1, 3].into_array().as_ref(), - ) - .unwrap(); + let taken = ree_array().take(buffer![9, 8, 1, 3].into_array()).unwrap(); let expected = PrimitiveArray::from_iter(vec![5i32, 5, 1, 4]).into_array(); assert_arrays_eq!(taken, expected); } #[test] fn ree_take_end() { - let taken = take(ree_array().as_ref(), buffer![11].into_array().as_ref()).unwrap(); + let taken = ree_array().take(buffer![11].into_array()).unwrap(); let expected = PrimitiveArray::from_iter(vec![5i32]).into_array(); assert_arrays_eq!(taken, expected); } @@ -129,7 +124,8 @@ mod test { #[test] #[should_panic] fn ree_take_out_of_bounds() { - let _array = take(ree_array().as_ref(), buffer![12].into_array().as_ref()) + let _array = ree_array() + .take(buffer![12].into_array()) .unwrap() .execute::(&mut LEGACY_SESSION.create_execution_ctx()) .unwrap(); @@ -138,7 +134,7 @@ mod test { #[test] fn sliced_take() { let sliced = ree_array().slice(4..9).unwrap(); - let taken = take(sliced.as_ref(), buffer![1, 3, 4].into_array().as_ref()).unwrap(); + let taken = sliced.take(buffer![1, 3, 4].into_array()).unwrap(); let expected = PrimitiveArray::from_iter(vec![4i32, 2, 5]).into_array(); assert_arrays_eq!(taken, expected); @@ -146,11 +142,9 @@ mod test { #[test] fn ree_take_nullable() { - let taken = take( - ree_array().as_ref(), - PrimitiveArray::from_option_iter([Some(1), None]).as_ref(), - ) - .unwrap(); + let taken = ree_array() + .take(PrimitiveArray::from_option_iter([Some(1), None]).to_array()) + .unwrap(); let expected = PrimitiveArray::from_option_iter([Some(1i32), None]); assert_arrays_eq!(taken, expected.to_array()); diff --git a/encodings/sequence/src/compute/take.rs b/encodings/sequence/src/compute/take.rs index 312be5baf44..96c002d1fd5 100644 --- a/encodings/sequence/src/compute/take.rs +++ b/encodings/sequence/src/compute/take.rs @@ -107,7 +107,6 @@ mod test { use vortex_array::Canonical; use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; - use vortex_array::compute::take; use vortex_dtype::Nullability; use crate::SequenceArray; @@ -171,7 +170,8 @@ mod test { fn test_bounds_check() { let array = SequenceArray::typed_new(0i32, 1i32, Nullability::NonNullable, 10).unwrap(); let indices = vortex_array::arrays::PrimitiveArray::from_iter([0i32, 20]); - let _array = take(array.as_ref(), indices.as_ref()) + let _array = array + .take(indices.to_array()) .unwrap() .execute::(&mut LEGACY_SESSION.create_execution_ctx()) .unwrap(); diff --git a/encodings/sparse/src/compute/take.rs b/encodings/sparse/src/compute/take.rs index 3ff7d74b0df..05c88ed787d 100644 --- a/encodings/sparse/src/compute/take.rs +++ b/encodings/sparse/src/compute/take.rs @@ -63,7 +63,6 @@ mod test { use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; - use vortex_array::compute::take; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_dtype::Nullability; @@ -91,7 +90,7 @@ mod test { fn take_with_non_zero_offset() { let sparse = sparse_array(); let sparse = sparse.slice(30..40).unwrap(); - let taken = take(&sparse, &buffer![6, 7, 8].into_array()).unwrap(); + let taken = sparse.take(buffer![6, 7, 8].into_array()).unwrap(); let expected = PrimitiveArray::from_option_iter([Option::::None, Some(0.47), None]); assert_arrays_eq!(taken, expected.to_array()); } @@ -99,7 +98,7 @@ mod test { #[test] fn sparse_take() { let sparse = sparse_array(); - let taken = take(&sparse, &buffer![0, 47, 47, 0, 99].into_array()).unwrap(); + let taken = sparse.take(buffer![0, 47, 47, 0, 99].into_array()).unwrap(); let expected = PrimitiveArray::from_option_iter([ Some(1.23f64), Some(9.99), @@ -113,7 +112,7 @@ mod test { #[test] fn nonexistent_take() { let sparse = sparse_array(); - let taken = take(&sparse, &buffer![69].into_array()).unwrap(); + let taken = sparse.take(buffer![69].into_array()).unwrap(); let expected = ConstantArray::new(test_array_fill_value(), 1).into_array(); assert_arrays_eq!(taken, expected); } @@ -122,7 +121,7 @@ mod test { fn ordered_take() { let sparse = sparse_array(); // Note: take returns a canonical array, not SparseArray - let taken = take(&sparse, &buffer![69, 37].into_array()).unwrap(); + let taken = sparse.take(buffer![69, 37].into_array()).unwrap(); // Index 69 is not in sparse array (fill value is null), index 37 has value 0.47 let expected = PrimitiveArray::from_option_iter([Option::::None, Some(0.47f64)]); assert_arrays_eq!(taken, expected.to_array()); @@ -138,12 +137,12 @@ mod test { ) .unwrap(); - let taken = take( - arr.as_ref(), - PrimitiveArray::from_option_iter([Some(2u32), Some(1u32), Option::::None]) - .as_ref(), - ) - .unwrap(); + let taken = arr + .take( + PrimitiveArray::from_option_iter([Some(2u32), Some(1u32), Option::::None]) + .to_array(), + ) + .unwrap(); let expected = PrimitiveArray::from_option_iter([Some(1), Some(10), Option::::None]); assert_arrays_eq!(taken, expected.to_array()); @@ -159,12 +158,12 @@ mod test { ) .unwrap(); - let taken = take( - arr.as_ref(), - PrimitiveArray::from_option_iter([Some(2u32), Some(1u32), Option::::None]) - .as_ref(), - ) - .unwrap(); + let taken = arr + .take( + PrimitiveArray::from_option_iter([Some(2u32), Some(1u32), Option::::None]) + .to_array(), + ) + .unwrap(); let expected = PrimitiveArray::from_option_iter([Some(1), Some(10), Option::::None]); assert_arrays_eq!(taken, expected.to_array()); diff --git a/encodings/zigzag/src/compute/mod.rs b/encodings/zigzag/src/compute/mod.rs index 42b9ea2d120..562af093215 100644 --- a/encodings/zigzag/src/compute/mod.rs +++ b/encodings/zigzag/src/compute/mod.rs @@ -77,7 +77,6 @@ mod tests { use vortex_array::assert_arrays_eq; use vortex_array::compute::conformance::binary_numeric::test_binary_numeric_array; use vortex_array::compute::conformance::consistency::test_array_consistency; - use vortex_array::compute::take; use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; use vortex_buffer::buffer; @@ -109,7 +108,7 @@ mod tests { ))?; let indices = buffer![0, 2].into_array(); - let actual = take(zigzag.as_ref(), &indices).unwrap(); + let actual = zigzag.take(indices.to_array()).unwrap(); let expected = zigzag_encode(PrimitiveArray::new(buffer![-189, 1], Validity::AllValid))?; assert_arrays_eq!(actual, expected); Ok(()) diff --git a/fuzz/src/array/mod.rs b/fuzz/src/array/mod.rs index dc9488829b1..52c098fbe81 100644 --- a/fuzz/src/array/mod.rs +++ b/fuzz/src/array/mod.rs @@ -557,8 +557,6 @@ pub fn run_fuzz_action(fuzz_action: FuzzArrayAction) -> crate::error::VortexFuzz use vortex_array::compute::mask; use vortex_array::compute::min_max; use vortex_array::compute::sum; - use vortex_array::compute::take; - let FuzzArrayAction { array, actions } = fuzz_action; let mut current_array = array.to_array(); @@ -581,7 +579,8 @@ pub fn run_fuzz_action(fuzz_action: FuzzArrayAction) -> crate::error::VortexFuzz if indices.is_empty() { return Ok(false); // Reject } - current_array = take(¤t_array, &indices) + current_array = current_array + .take(indices) .vortex_expect("take operation should succeed in fuzz test"); assert_array_eq(&expected.array(), ¤t_array, i)?; } diff --git a/vortex-array/benches/take_fsl.rs b/vortex-array/benches/take_fsl.rs index ef8d3e00d87..cd5a7233196 100644 --- a/vortex-array/benches/take_fsl.rs +++ b/vortex-array/benches/take_fsl.rs @@ -14,9 +14,9 @@ use divan::Bencher; use rand::Rng; use rand::SeedableRng; use rand::rngs::StdRng; +use vortex_array::Array; use vortex_array::IntoArray; use vortex_array::arrays::FixedSizeListArray; -use vortex_array::compute::take; use vortex_array::validity::Validity; use vortex_buffer::Buffer; @@ -61,7 +61,7 @@ fn take_fsl_random(bencher: Bencher, num_indices: usize) bencher .with_inputs(|| (&fsl, &indices_array)) - .bench_refs(|(array, indices)| take(array.as_ref(), indices.as_ref()).unwrap()); + .bench_refs(|(array, indices)| array.take(indices.to_array()).unwrap()); } #[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] @@ -80,5 +80,5 @@ fn take_fsl_nullable_random(bencher: Bencher, num_indice bencher .with_inputs(|| (&fsl, &indices_array)) - .bench_refs(|(array, indices)| take(array.as_ref(), indices.as_ref()).unwrap()); + .bench_refs(|(array, indices)| array.take(indices.to_array()).unwrap()); } diff --git a/vortex-array/benches/take_struct.rs b/vortex-array/benches/take_struct.rs index ff4c58276e0..dc1b6547b0b 100644 --- a/vortex-array/benches/take_struct.rs +++ b/vortex-array/benches/take_struct.rs @@ -8,9 +8,9 @@ use rand::Rng; use rand::SeedableRng; use rand::distr::Uniform; use rand::rngs::StdRng; +use vortex_array::Array; use vortex_array::IntoArray; use vortex_array::arrays::StructArray; -use vortex_array::compute::take; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_dtype::FieldNames; @@ -48,7 +48,7 @@ fn take_struct_simple(bencher: Bencher) { bencher .with_inputs(|| (&struct_array, &indices_array)) - .bench_refs(|(array, indices)| take(array.as_ref(), indices.as_ref()).unwrap()); + .bench_refs(|(array, indices)| array.take(indices.to_array()).unwrap()); } #[divan::bench(args = [8])] @@ -79,7 +79,7 @@ fn take_struct_wide(bencher: Bencher, width: usize) { bencher .with_inputs(|| (&struct_array, &indices_array)) - .bench_refs(|(array, indices)| take(array.as_ref(), indices.as_ref()).unwrap()); + .bench_refs(|(array, indices)| array.take(indices.to_array()).unwrap()); } #[divan::bench] @@ -107,5 +107,5 @@ fn take_struct_sequential_indices(bencher: Bencher) { bencher .with_inputs(|| (&struct_array, &indices_array)) - .bench_refs(|(array, indices)| take(array.as_ref(), indices.as_ref()).unwrap()); + .bench_refs(|(array, indices)| array.take(indices.to_array()).unwrap()); } diff --git a/vortex-array/benches/varbinview_compact.rs b/vortex-array/benches/varbinview_compact.rs index 1b1be44cf48..3ad3db497da 100644 --- a/vortex-array/benches/varbinview_compact.rs +++ b/vortex-array/benches/varbinview_compact.rs @@ -10,7 +10,6 @@ use vortex_array::IntoArray; use vortex_array::ToCanonical; use vortex_array::arrays::VarBinViewArray; use vortex_array::builders::VarBinViewBuilder; -use vortex_array::compute::take; use vortex_buffer::Buffer; use vortex_dtype::DType; use vortex_dtype::Nullability; @@ -42,8 +41,9 @@ fn compact_impl(bencher: Bencher, (output_size, utilization_pct): (usize, usize) let base_size = (output_size * 100) / utilization_pct; let base_array = build_varbinview_fixture(base_size); let indices = random_indices(output_size, base_size); - let taken = - take(base_array.as_ref(), &indices).vortex_expect("operation should succeed in benchmark"); + let taken = base_array + .take(indices) + .vortex_expect("operation should succeed in benchmark"); let array = taken.to_varbinview(); bencher.with_inputs(|| &array).bench_refs(|array| { diff --git a/vortex-array/src/arrays/bool/compute/take.rs b/vortex-array/src/arrays/bool/compute/take.rs index e9fb6bd44dc..95369b12024 100644 --- a/vortex-array/src/arrays/bool/compute/take.rs +++ b/vortex-array/src/arrays/bool/compute/take.rs @@ -88,7 +88,6 @@ mod test { use crate::arrays::primitive::PrimitiveArray; use crate::assert_arrays_eq; use crate::compute::conformance::take::test_take_conformance; - use crate::compute::take; use crate::validity::Validity; #[test] @@ -101,7 +100,8 @@ mod test { Some(false), ]); - let b = take(reference.as_ref(), buffer![0, 3, 4].into_array().as_ref()) + let b = reference + .take(buffer![0, 3, 4].into_array()) .unwrap() .to_bool(); assert_eq!( @@ -110,7 +110,7 @@ mod test { ); let all_invalid_indices = PrimitiveArray::from_option_iter([None::, None, None]); - let b = take(reference.as_ref(), all_invalid_indices.as_ref()).unwrap(); + let b = reference.take(all_invalid_indices.to_array()).unwrap(); assert_arrays_eq!(b, BoolArray::from_iter([None, None, None])); } @@ -121,7 +121,7 @@ mod test { buffer![0, 3, 100], Validity::Array(BoolArray::from_iter([true, true, false]).to_array()), ); - let actual = take(values.as_ref(), indices.as_ref()).unwrap(); + let actual = values.take(indices.to_array()).unwrap(); // position 3 is null, the third index is null assert_arrays_eq!(actual, BoolArray::from_iter([Some(false), None, None])); @@ -134,7 +134,7 @@ mod test { buffer![0, 3, 100], Validity::Array(BoolArray::from_iter([true, true, false]).to_array()), ); - let actual = take(values.as_ref(), indices.as_ref()).unwrap(); + let actual = values.take(indices.to_array()).unwrap(); // the third index is null assert_arrays_eq!( actual, @@ -149,7 +149,7 @@ mod test { buffer![0, 3, 100], Validity::Array(BoolArray::from_iter([false, false, false]).to_array()), ); - let actual = take(values.as_ref(), indices.as_ref()).unwrap(); + let actual = values.take(indices.to_array()).unwrap(); assert_arrays_eq!(actual, BoolArray::from_iter([None, None, None])); } @@ -160,7 +160,7 @@ mod test { buffer![0, 3, 100], Validity::Array(BoolArray::from_iter([false, false, false]).to_array()), ); - let actual = take(values.as_ref(), indices.as_ref()).unwrap(); + let actual = values.take(indices.to_array()).unwrap(); assert_arrays_eq!(actual, BoolArray::from_iter([None, None, None])); } diff --git a/vortex-array/src/arrays/chunked/compute/take.rs b/vortex-array/src/arrays/chunked/compute/take.rs index 4210e764fb3..a56868c6ed0 100644 --- a/vortex-array/src/arrays/chunked/compute/take.rs +++ b/vortex-array/src/arrays/chunked/compute/take.rs @@ -106,7 +106,6 @@ mod test { use crate::arrays::chunked::ChunkedArray; use crate::assert_arrays_eq; use crate::compute::conformance::take::test_take_conformance; - use crate::compute::take; use crate::validity::Validity; #[test] @@ -118,7 +117,7 @@ mod test { assert_eq!(arr.len(), 9); let indices = buffer![0u64, 0, 6, 4].into_array(); - let result = take(arr.as_ref(), indices.as_ref()).unwrap(); + let result = arr.take(indices.to_array()).unwrap(); assert_arrays_eq!(result, PrimitiveArray::from_iter([1i32, 1, 1, 2])); } @@ -130,11 +129,9 @@ mod test { let arr = ChunkedArray::from_iter(vec![struct_array.to_array(), struct_array.to_array()]); - let result = take( - arr.as_ref(), - PrimitiveArray::from_option_iter(vec![Some(0), None, Some(101)]).as_ref(), - ) - .unwrap(); + let result = arr + .take(PrimitiveArray::from_option_iter(vec![Some(0), None, Some(101)]).to_array()) + .unwrap(); let expect = StructArray::try_new( FieldNames::default(), @@ -155,7 +152,7 @@ mod test { assert_eq!(arr.len(), 9); let indices = PrimitiveArray::empty::(Nullability::NonNullable); - let result = take(arr.as_ref(), indices.as_ref()).unwrap(); + let result = arr.take(indices.to_array()).unwrap(); assert!(result.is_empty()); assert_eq!(result.dtype(), arr.dtype()); diff --git a/vortex-array/src/arrays/constant/compute/take.rs b/vortex-array/src/arrays/constant/compute/take.rs index 44a00f76ca4..3f565a77729 100644 --- a/vortex-array/src/arrays/constant/compute/take.rs +++ b/vortex-array/src/arrays/constant/compute/take.rs @@ -72,21 +72,20 @@ mod tests { use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::compute::conformance::take::test_take_conformance; - use crate::compute::take; use crate::validity::Validity; #[test] fn take_nullable_indices() { let array = ConstantArray::new(42, 10).to_array(); - let taken = take( - &array, - &PrimitiveArray::new( - buffer![0, 5, 7], - Validity::from_iter(vec![false, true, false]), + let taken = array + .take( + PrimitiveArray::new( + buffer![0, 5, 7], + Validity::from_iter(vec![false, true, false]), + ) + .into_array(), ) - .into_array(), - ) - .unwrap(); + .unwrap(); let valid_indices: &[usize] = &[1usize]; assert_eq!( &array.dtype().with_nullability(Nullability::Nullable), @@ -108,11 +107,9 @@ mod tests { #[test] fn take_all_valid_indices() { let array = ConstantArray::new(42, 10).to_array(); - let taken = take( - &array, - &PrimitiveArray::new(buffer![0, 5, 7], Validity::AllValid).into_array(), - ) - .unwrap(); + let taken = array + .take(PrimitiveArray::new(buffer![0, 5, 7], Validity::AllValid).into_array()) + .unwrap(); assert_eq!( &array.dtype().with_nullability(Nullability::Nullable), taken.dtype() diff --git a/vortex-array/src/arrays/decimal/compute/take.rs b/vortex-array/src/arrays/decimal/compute/take.rs index 2e6eb4abfb2..29869c2b6c0 100644 --- a/vortex-array/src/arrays/decimal/compute/take.rs +++ b/vortex-array/src/arrays/decimal/compute/take.rs @@ -59,7 +59,6 @@ mod tests { use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::compute::conformance::take::test_take_conformance; - use crate::compute::take; use crate::validity::Validity; #[test] @@ -72,7 +71,7 @@ mod tests { ); let indices = buffer![0, 2, 3].into_array(); - let taken = take(array.as_ref(), indices.as_ref()).unwrap(); + let taken = array.take(indices.to_array()).unwrap(); let expected = DecimalArray::from_iter([10i128, 12, 13], ddtype); assert_arrays_eq!(expected, taken); @@ -88,7 +87,7 @@ mod tests { ); let indices = PrimitiveArray::from_option_iter([None, Some(2), Some(3)]).into_array(); - let taken = take(array.as_ref(), indices.as_ref()).unwrap(); + let taken = array.take(indices.to_array()).unwrap(); let expected = DecimalArray::from_option_iter([None, Some(12i128), Some(13)], ddtype); assert_arrays_eq!(expected, taken); diff --git a/vortex-array/src/arrays/dict/compute/mod.rs b/vortex-array/src/arrays/dict/compute/mod.rs index fe8b053a53f..cf5fe445857 100644 --- a/vortex-array/src/arrays/dict/compute/mod.rs +++ b/vortex-array/src/arrays/dict/compute/mod.rs @@ -80,8 +80,6 @@ mod test { use crate::compute::conformance::filter::test_filter_conformance; use crate::compute::conformance::mask::test_mask_conformance; use crate::compute::conformance::take::test_take_conformance; - use crate::compute::take; - #[test] fn canonicalise_nullable_primitive() { let values: Vec> = (0..65) @@ -226,12 +224,10 @@ mod test { let array = dict_encode(buffer![1, 2].into_array().as_ref()).unwrap(); assert_eq!( - take( - array.as_ref(), - PrimitiveArray::from_option_iter([Option::::None]).as_ref() - ) - .unwrap() - .dtype(), + array + .take(PrimitiveArray::from_option_iter([Option::::None]).to_array()) + .unwrap() + .dtype(), &DType::Primitive(I32, Nullability::Nullable) ); } diff --git a/vortex-array/src/arrays/fixed_size_list/tests/take.rs b/vortex-array/src/arrays/fixed_size_list/tests/take.rs index 00795ac4516..defbc56758b 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/take.rs @@ -21,7 +21,6 @@ use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; use crate::builders::FixedSizeListBuilder; use crate::compute::conformance::take::test_take_conformance; -use crate::compute::take; use crate::validity::Validity; // Conformance tests for common take scenarios. @@ -43,7 +42,7 @@ fn test_take_basic_smoke_test() { let fsl = FixedSizeListArray::new(elements.into_array(), 2, Validity::NonNullable, 3); let indices = buffer![2u32, 0, 1].into_array(); - let result = take(fsl.as_ref(), indices.as_ref()).unwrap(); + let result = fsl.take(indices.to_array()).unwrap(); // Expected: [[5,6], [1,2], [3,4]] let expected = FixedSizeListArray::new( @@ -90,7 +89,7 @@ fn test_take_degenerate_lists( // Also test the specific behavior. let indices_array = PrimitiveArray::from_option_iter(indices); - let result = take(fsl.as_ref(), indices_array.as_ref()).unwrap(); + let result = fsl.take(indices_array.to_array()).unwrap(); assert_eq!(result.len(), expected_len); for (i, expected_null) in expected_nulls.iter().enumerate() { @@ -104,7 +103,7 @@ fn test_take_large_list_size() { let fsl = FixedSizeListArray::new(elements, 100, Validity::NonNullable, 3); let indices = buffer![2u16, 0].into_array(); - let result = take(fsl.as_ref(), indices.as_ref()).unwrap(); + let result = fsl.take(indices.to_array()).unwrap(); // Expected: [[200..300], [0..100]] let expected_elems = PrimitiveArray::from_iter((200i32..300).chain(0..100)).into_array(); @@ -119,7 +118,7 @@ fn test_take_fsl_with_null_indices_preserves_elements() { // Indices with nulls: [1, null, 0]. let indices = PrimitiveArray::from_option_iter([Some(1u32), None, Some(0)]); - let result = take(fsl.as_ref(), indices.as_ref()).unwrap(); + let result = fsl.take(indices.to_array()).unwrap(); // Expected: [[3,4], null, [1,2]] let expected = FixedSizeListArray::new( @@ -185,7 +184,7 @@ fn test_take_nullable_arrays_fsl_specific( // Create indices (with possible nulls). let indices_array = PrimitiveArray::from_option_iter(indices.clone()); - let result = take(fsl.as_ref(), indices_array.as_ref()).unwrap(); + let result = fsl.take(indices_array.to_array()).unwrap(); assert_eq!(result.len(), indices.len()); for (i, expected_null) in expected_nulls.iter().enumerate() { diff --git a/vortex-array/src/arrays/list/compute/take.rs b/vortex-array/src/arrays/list/compute/take.rs index 62fd7b1cf8d..df75b7f72c7 100644 --- a/vortex-array/src/arrays/list/compute/take.rs +++ b/vortex-array/src/arrays/list/compute/take.rs @@ -196,7 +196,6 @@ mod test { use crate::arrays::PrimitiveArray; use crate::arrays::list::ListArray; use crate::compute::conformance::take::test_take_conformance; - use crate::compute::take; use crate::validity::Validity; #[test] @@ -212,7 +211,7 @@ mod test { let idx = PrimitiveArray::from_option_iter(vec![Some(0), None, Some(1), Some(3)]).to_array(); - let result = take(&list, &idx).unwrap(); + let result = list.take(idx.to_array()).unwrap(); assert_eq!( result.dtype(), @@ -270,7 +269,7 @@ mod test { let idx = PrimitiveArray::from_option_iter(vec![Some(0), Some(1), None]).to_array(); // since idx is nullable, the final list will also be nullable - let result = take(&list, &idx).unwrap(); + let result = list.take(idx.to_array()).unwrap(); assert_eq!( result.dtype(), &DType::List( @@ -292,7 +291,7 @@ mod test { let idx = buffer![1, 0, 2].into_array(); - let result = take(&list, &idx).unwrap(); + let result = list.take(idx.to_array()).unwrap(); assert_eq!( result.dtype(), @@ -347,7 +346,7 @@ mod test { let idx = PrimitiveArray::empty::(Nullability::Nullable).to_array(); - let result = take(&list, &idx).unwrap(); + let result = list.take(idx.to_array()).unwrap(); assert_eq!( result.dtype(), &DType::List( @@ -410,7 +409,7 @@ mod test { // Take the same large list twice - would overflow u8 but works with u64. let idx = buffer![0u8, 0].into_array(); - let result = take(&list, &idx).unwrap(); + let result = list.take(idx.to_array()).unwrap(); assert_eq!(result.len(), 2); @@ -431,7 +430,7 @@ mod test { // Take the same large list twice - would overflow u8 but works with u64. let idx = PrimitiveArray::from_option_iter(vec![Some(0u8), None, Some(0u8)]).to_array(); - let result = take(&list, &idx).unwrap(); + let result = list.take(idx.to_array()).unwrap(); assert_eq!(result.len(), 3); @@ -461,7 +460,7 @@ mod test { let idx = buffer![0u32, 1, 0, 1].into_array(); // This should not panic - result should have length 4. - let result = take(&list, &idx).unwrap(); + let result = list.take(idx.to_array()).unwrap(); assert_eq!(result.len(), 4); } } diff --git a/vortex-array/src/arrays/listview/tests/take.rs b/vortex-array/src/arrays/listview/tests/take.rs index 9c5ec621e20..7f596d25847 100644 --- a/vortex-array/src/arrays/listview/tests/take.rs +++ b/vortex-array/src/arrays/listview/tests/take.rs @@ -17,7 +17,6 @@ use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::compute::conformance::take::test_take_conformance; -use crate::compute::take; use crate::validity::Validity; // Conformance tests for common take scenarios. @@ -47,7 +46,7 @@ fn test_take_preserves_unreferenced_elements() { // Take only 2 lists. let indices = buffer![1u32, 3].into_array(); - let result = take(&listview, &indices).unwrap(); + let result = listview.take(indices.to_array()).unwrap(); let result_list = result.to_listview(); assert_eq!(result_list.len(), 2); @@ -76,7 +75,7 @@ fn test_take_with_gaps() { ListViewArray::new(elements.clone(), offsets, sizes, Validity::NonNullable).to_array(); let indices = buffer![1u32, 3, 4, 2].into_array(); - let result = take(&listview, &indices).unwrap(); + let result = listview.take(indices.to_array()).unwrap(); let result_list = result.to_listview(); // Verify the entire elements array is preserved including gaps. @@ -111,7 +110,7 @@ fn test_take_constant_arrays() { .to_array(); let indices = buffer![3u32, 0, 2].into_array(); - let result = take(&const_offset_list, &indices).unwrap(); + let result = const_offset_list.take(indices.to_array()).unwrap(); let result_list = result.to_listview(); assert_eq!(result_list.len(), 3); @@ -135,7 +134,7 @@ fn test_take_constant_arrays() { .to_array(); let indices2 = buffer![2u32, 0].into_array(); - let result2 = take(&both_const_list, &indices2).unwrap(); + let result2 = both_const_list.take(indices2.to_array()).unwrap(); let result2_list = result2.to_listview(); assert_eq!(result2_list.len(), 2); @@ -161,7 +160,7 @@ fn test_take_extreme_offsets() { // Take only 2 lists, demonstrating we keep all 10000 elements. let indices = buffer![1u32, 4].into_array(); - let result = take(&listview, &indices).unwrap(); + let result = listview.take(indices.to_array()).unwrap(); let result_list = result.to_listview(); assert_eq!(result_list.len(), 2); diff --git a/vortex-array/src/arrays/null/compute/mod.rs b/vortex-array/src/arrays/null/compute/mod.rs index 620d08ad19a..2e87e3df672 100644 --- a/vortex-array/src/arrays/null/compute/mod.rs +++ b/vortex-array/src/arrays/null/compute/mod.rs @@ -23,7 +23,6 @@ mod test { use crate::compute::conformance::filter::test_filter_conformance; use crate::compute::conformance::mask::test_mask_conformance; use crate::compute::conformance::take::test_take_conformance; - use crate::compute::take; #[test] fn test_slice_nulls() { @@ -37,7 +36,8 @@ mod test { #[test] fn test_take_nulls() { let nulls = NullArray::new(10); - let taken = take(nulls.as_ref(), &buffer![0u64, 2, 4, 6, 8].into_array()) + let taken = nulls + .take(buffer![0u64, 2, 4, 6, 8].into_array()) .unwrap() .to_null(); diff --git a/vortex-array/src/arrays/primitive/compute/take/mod.rs b/vortex-array/src/arrays/primitive/compute/take/mod.rs index be832f50cb6..26457aa1af6 100644 --- a/vortex-array/src/arrays/primitive/compute/take/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/take/mod.rs @@ -126,7 +126,6 @@ mod test { use crate::arrays::PrimitiveArray; use crate::arrays::primitive::compute::take::take_primitive_scalar; use crate::compute::conformance::take::test_take_conformance; - use crate::compute::take; use crate::validity::Validity; #[test] @@ -146,7 +145,7 @@ mod test { buffer![0, 3, 4], Validity::Array(BoolArray::from_iter([true, true, false]).into_array()), ); - let actual = take(values.as_ref(), indices.as_ref()).unwrap(); + let actual = values.take(indices.to_array()).unwrap(); assert_eq!( actual.scalar_at(0).vortex_expect("no fail"), Scalar::from(Some(1)) diff --git a/vortex-array/src/arrays/struct_/compute/mod.rs b/vortex-array/src/arrays/struct_/compute/mod.rs index 3694beadfec..a2d9f5cd4fd 100644 --- a/vortex-array/src/arrays/struct_/compute/mod.rs +++ b/vortex-array/src/arrays/struct_/compute/mod.rs @@ -25,7 +25,10 @@ mod tests { use vortex_error::VortexExpect; use crate::Array; + use crate::Canonical; use crate::IntoArray as _; + use crate::LEGACY_SESSION; + use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; @@ -36,7 +39,6 @@ mod tests { use crate::compute::conformance::mask::test_mask_conformance; use crate::compute::conformance::take::test_take_conformance; use crate::compute::is_constant; - use crate::compute::take; use crate::validity::Validity; #[test] @@ -44,7 +46,7 @@ mod tests { let struct_arr = StructArray::try_new(FieldNames::empty(), vec![], 10, Validity::NonNullable).unwrap(); let indices = PrimitiveArray::from_option_iter([Some(1), None]); - let taken = take(struct_arr.as_ref(), indices.as_ref()).unwrap(); + let taken = struct_arr.take(indices.to_array()).unwrap(); assert_arrays_eq!( taken, @@ -65,16 +67,24 @@ mod tests { ) .unwrap(); let indices = PrimitiveArray::from_option_iter([Option::::None]); - let taken = take(struct_arr.as_ref(), indices.as_ref()).unwrap(); + let taken = struct_arr + .take(indices.to_array()) + .unwrap() + .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .unwrap(); assert_eq!(taken.len(), 1); - assert!(taken.all_invalid().unwrap()); + assert!(taken.into_array().all_invalid().unwrap()); } #[test] fn take_empty_primitive_with_nullable_indices() { let arr = PrimitiveArray::from_iter(Vec::::new()); let indices = PrimitiveArray::from_option_iter([Option::::None]); - let taken = take(arr.as_ref(), indices.as_ref()).unwrap(); + let taken = arr + .take(indices.to_array()) + .unwrap() + .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .unwrap(); assert_eq!(taken.len(), 1); } @@ -84,7 +94,7 @@ mod tests { StructArray::from_fields(&[("a", PrimitiveArray::from_iter(0..10).to_array())]) .unwrap(); let indices = PrimitiveArray::from_option_iter([Some(1), None]); - let taken = take(struct_arr.as_ref(), indices.as_ref()).unwrap(); + let taken = struct_arr.take(indices.to_array()).unwrap(); assert_arrays_eq!( taken, StructArray::try_from_iter_with_validity( diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index 83dec64b6bb..143b1065450 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -261,7 +261,6 @@ mod tests { use crate::arrays::VarBinViewArray; use crate::assert_arrays_eq; use crate::compute::conformance::take::test_take_conformance; - use crate::compute::take; use crate::validity::Validity; #[test] @@ -271,14 +270,14 @@ mod tests { let idx1: PrimitiveArray = (0..1).collect(); assert_eq!( - take(arr.as_ref(), idx1.as_ref()).unwrap().dtype(), + arr.take(idx1.to_array()).unwrap().dtype(), &DType::Utf8(Nullability::NonNullable) ); let idx2: PrimitiveArray = PrimitiveArray::from_option_iter(vec![Some(0)]); assert_eq!( - take(arr.as_ref(), idx2.as_ref()).unwrap().dtype(), + arr.take(idx2.to_array()).unwrap().dtype(), &DType::Utf8(Nullability::Nullable) ); } @@ -315,7 +314,7 @@ mod tests { ); let indices = buffer![0u32; 3].into_array(); - let taken = take(array.as_ref(), indices.as_ref()).unwrap(); + let taken = array.take(indices.to_array()).unwrap(); let expected = VarBinViewArray::from_iter( [Some(scream.clone()), Some(scream.clone()), Some(scream)], diff --git a/vortex-array/src/arrays/varbinview/compact.rs b/vortex-array/src/arrays/varbinview/compact.rs index 58c0ea4bb11..bd9b2d236f2 100644 --- a/vortex-array/src/arrays/varbinview/compact.rs +++ b/vortex-array/src/arrays/varbinview/compact.rs @@ -194,8 +194,6 @@ mod tests { use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; use crate::assert_arrays_eq; - use crate::compute::take; - #[test] fn test_optimize_compacts_buffers() { // Create a VarBinViewArray with some long strings that will create multiple buffers @@ -213,7 +211,7 @@ mod tests { // Take only the first and last elements (indices 0 and 4) let indices = buffer![0u32, 4u32].into_array(); - let taken = take(original.as_ref(), &indices).unwrap(); + let taken = original.take(indices.to_array()).unwrap(); let taken = taken .execute::(&mut LEGACY_SESSION.create_execution_ctx()) .unwrap(); @@ -252,7 +250,7 @@ mod tests { // Take only the first and third long strings (indices 0 and 2) let indices = buffer![0u32, 2u32].into_array(); - let taken = take(original.as_ref(), &indices).unwrap(); + let taken = original.take(indices.to_array()).unwrap(); let taken_array = taken .execute::(&mut LEGACY_SESSION.create_execution_ctx()) .unwrap(); @@ -318,7 +316,7 @@ mod tests { // Take only first element let indices = buffer![0u32].into_array(); - let taken = take(original.as_ref(), &indices).unwrap(); + let taken = original.take(indices.to_array()).unwrap(); let taken = taken .execute::(&mut LEGACY_SESSION.create_execution_ctx()) .unwrap(); @@ -343,7 +341,7 @@ mod tests { // Take only first and last elements let indices = buffer![0u32, 2u32].into_array(); - let taken = take(original.as_ref(), &indices).unwrap(); + let taken = original.take(indices.to_array()).unwrap(); let taken = taken .clone() .execute::(&mut LEGACY_SESSION.create_execution_ctx()) @@ -399,7 +397,7 @@ mod tests { // Take every other element to create mixed utilization let indices_array = buffer![0u32, 2u32, 4u32, 6u32, 8u32].into_array(); - let taken = take(original.as_ref(), &indices_array).unwrap(); + let taken = original.take(indices_array.to_array()).unwrap(); let taken = taken .execute::(&mut LEGACY_SESSION.create_execution_ctx()) .unwrap(); @@ -425,7 +423,7 @@ mod tests { // Take only the first 5 elements - they should be in a contiguous range at the start let indices_array = buffer![0u32, 1u32, 2u32, 3u32, 4u32].into_array(); - let taken = take(original.as_ref(), &indices_array).unwrap(); + let taken = original.take(indices_array.to_array()).unwrap(); let taken = taken .execute::(&mut LEGACY_SESSION.create_execution_ctx()) .unwrap(); diff --git a/vortex-array/src/arrays/varbinview/compute/mod.rs b/vortex-array/src/arrays/varbinview/compute/mod.rs index f538e5c9069..f97f7774179 100644 --- a/vortex-array/src/arrays/varbinview/compute/mod.rs +++ b/vortex-array/src/arrays/varbinview/compute/mod.rs @@ -20,8 +20,6 @@ mod tests { use crate::array::Array; use crate::arrays::VarBinViewArray; use crate::canonical::ToCanonical; - use crate::compute::take; - #[test] fn take_nullable() { let arr = VarBinViewArray::from_iter_nullable_str([ @@ -33,7 +31,7 @@ mod tests { Some("six"), ]); - let taken = take(arr.as_ref(), &buffer![0, 3].into_array()).unwrap(); + let taken = arr.take(buffer![0, 3].into_array()).unwrap(); assert!(taken.dtype().is_nullable()); assert_eq!( diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index 48dc08db32f..55e12dff324 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -97,7 +97,6 @@ mod tests { use crate::arrays::VarBinViewArray; use crate::canonical::ToCanonical; use crate::compute::conformance::take::test_take_conformance; - use crate::compute::take; use crate::validity::Validity; #[test] @@ -111,7 +110,7 @@ mod tests { Some("six"), ]); - let taken = take(arr.as_ref(), &buffer![0, 3].into_array()).unwrap(); + let taken = arr.take(buffer![0, 3].into_array()).unwrap(); assert!(taken.dtype().is_nullable()); assert_eq!( @@ -132,7 +131,7 @@ mod tests { Validity::from(BitBuffer::from(vec![true, false])), ); - let taken = take(arr.as_ref(), indices.as_ref()).unwrap(); + let taken = arr.take(indices.to_array()).unwrap(); assert!(taken.dtype().is_nullable()); assert_eq!( diff --git a/vortex-array/src/compute/conformance/consistency.rs b/vortex-array/src/compute/conformance/consistency.rs index 26fb8076465..dd33ec8d69b 100644 --- a/vortex-array/src/compute/conformance/consistency.rs +++ b/vortex-array/src/compute/conformance/consistency.rs @@ -38,7 +38,6 @@ use crate::compute::compare; use crate::compute::invert; use crate::compute::mask; use crate::compute::or; -use crate::compute::take; /// Tests that filter and take operations produce consistent results. /// @@ -75,8 +74,9 @@ fn test_filter_take_consistency(array: &dyn Array) { let indices_array = PrimitiveArray::from_iter(indices).into_array(); // Take using those indices - let taken = - take(array, &indices_array).vortex_expect("take should succeed in conformance test"); + let taken = array + .take(indices_array.to_array()) + .vortex_expect("take should succeed in conformance test"); // Results should be identical assert_eq!( @@ -356,7 +356,9 @@ fn test_take_slice_consistency(array: &dyn Array) { // Take indices [1, 2, 3] let end = 4.min(len); let indices = PrimitiveArray::from_iter((1..end).map(|i| i as u64)).into_array(); - let taken = take(array, &indices).vortex_expect("take should succeed in conformance test"); + let taken = array + .take(indices.to_array()) + .vortex_expect("take should succeed in conformance test"); // Slice from 1 to end let sliced = array @@ -441,7 +443,9 @@ fn test_take_repeated_indices(array: &dyn Array) { // Take the first element three times let indices = PrimitiveArray::from_iter([0u64, 0, 0]).into_array(); - let taken = take(array, &indices).vortex_expect("take should succeed in conformance test"); + let taken = array + .take(indices.to_array()) + .vortex_expect("take should succeed in conformance test"); assert_eq!(taken.len(), 3); for i in 0..3 { @@ -506,8 +510,9 @@ fn test_empty_operations_consistency(array: &dyn Array) { // Empty take let empty_indices = PrimitiveArray::empty::(Nullability::NonNullable).into_array(); - let empty_take = - take(array, &empty_indices).vortex_expect("take should succeed in conformance test"); + let empty_take = array + .take(empty_indices.to_array()) + .vortex_expect("take should succeed in conformance test"); assert_eq!(empty_take.len(), 0); assert_eq!(empty_take.dtype(), array.dtype()); @@ -530,7 +535,9 @@ fn test_take_preserves_properties(array: &dyn Array) { // Take all elements in original order let indices = PrimitiveArray::from_iter((0..len).map(|i| i as u64)).into_array(); - let taken = take(array, &indices).vortex_expect("take should succeed in conformance test"); + let taken = array + .take(indices.to_array()) + .vortex_expect("take should succeed in conformance test"); // Should be identical to original assert_eq!(taken.len(), array.len()); @@ -573,7 +580,9 @@ fn test_nullable_indices_consistency(array: &dyn Array) { // Create nullable indices where some indices are null let indices = PrimitiveArray::from_option_iter([Some(0u64), None, Some(2u64)]).into_array(); - let taken = take(array, &indices).vortex_expect("take should succeed in conformance test"); + let taken = array + .take(indices.to_array()) + .vortex_expect("take should succeed in conformance test"); // Result should have nulls where indices were null assert_eq!( @@ -637,8 +646,9 @@ fn test_large_array_consistency(array: &dyn Array) { // Test with every 10th element let indices: Vec = (0..len).step_by(10).map(|i| i as u64).collect(); let indices_array = PrimitiveArray::from_iter(indices).into_array(); - let taken = - take(array, &indices_array).vortex_expect("take should succeed in conformance test"); + let taken = array + .take(indices_array.to_array()) + .vortex_expect("take should succeed in conformance test"); // Create equivalent filter mask let mask_pattern: Vec = (0..len).map(|i| i % 10 == 0).collect(); diff --git a/vortex-array/src/compute/conformance/take.rs b/vortex-array/src/compute/conformance/take.rs index 68f316df251..32a1251bbb2 100644 --- a/vortex-array/src/compute/conformance/take.rs +++ b/vortex-array/src/compute/conformance/take.rs @@ -9,7 +9,6 @@ use crate::Array; use crate::Canonical; use crate::IntoArray as _; use crate::arrays::PrimitiveArray; -use crate::compute::take; /// Test conformance of the take compute function for an array. /// @@ -55,8 +54,9 @@ pub fn test_take_conformance(array: &dyn Array) { fn test_take_all(array: &dyn Array) { let len = array.len(); let indices = PrimitiveArray::from_iter(0..len as u64); - let result = - take(array, indices.as_ref()).vortex_expect("take should succeed in conformance test"); + let result = array + .take(indices.to_array()) + .vortex_expect("take should succeed in conformance test"); assert_eq!(result.len(), len); assert_eq!(result.dtype(), array.dtype()); @@ -94,8 +94,9 @@ fn test_take_all(array: &dyn Array) { fn test_take_none(array: &dyn Array) { let indices: PrimitiveArray = PrimitiveArray::from_iter::<[u64; 0]>([]); - let result = - take(array, indices.as_ref()).vortex_expect("take should succeed in conformance test"); + let result = array + .take(indices.to_array()) + .vortex_expect("take should succeed in conformance test"); assert_eq!(result.len(), 0); assert_eq!(result.dtype(), array.dtype()); @@ -110,7 +111,8 @@ fn test_take_selective(array: &dyn Array) { let expected_len = indices.len(); let indices_array = PrimitiveArray::from_iter(indices.clone()); - let result = take(array, indices_array.as_ref()) + let result = array + .take(indices_array.to_array()) .vortex_expect("take should succeed in conformance test"); assert_eq!(result.len(), expected_len); @@ -130,8 +132,9 @@ fn test_take_selective(array: &dyn Array) { fn test_take_first_and_last(array: &dyn Array) { let len = array.len(); let indices = PrimitiveArray::from_iter([0u64, (len - 1) as u64]); - let result = - take(array, indices.as_ref()).vortex_expect("take should succeed in conformance test"); + let result = array + .take(indices.to_array()) + .vortex_expect("take should succeed in conformance test"); assert_eq!(result.len(), 2); assert_eq!( @@ -166,8 +169,9 @@ fn test_take_with_nullable_indices(array: &dyn Array) { }; let indices = PrimitiveArray::from_option_iter(indices_vec.clone()); - let result = - take(array, indices.as_ref()).vortex_expect("take should succeed in conformance test"); + let result = array + .take(indices.to_array()) + .vortex_expect("take should succeed in conformance test"); assert_eq!(result.len(), indices_vec.len()); assert_eq!( @@ -206,8 +210,9 @@ fn test_take_repeated_indices(array: &dyn Array) { // Take the first element multiple times let indices = buffer![0u64, 0, 0].into_array(); - let result = - take(array, indices.as_ref()).vortex_expect("take should succeed in conformance test"); + let result = array + .take(indices.to_array()) + .vortex_expect("take should succeed in conformance test"); assert_eq!(result.len(), 3); let first_elem = array @@ -225,8 +230,9 @@ fn test_take_repeated_indices(array: &dyn Array) { fn test_empty_indices(array: &dyn Array) { let indices = PrimitiveArray::empty::(Nullability::NonNullable); - let result = - take(array, indices.as_ref()).vortex_expect("take should succeed in conformance test"); + let result = array + .take(indices.to_array()) + .vortex_expect("take should succeed in conformance test"); assert_eq!(result.len(), 0); assert_eq!(result.dtype(), array.dtype()); @@ -236,8 +242,9 @@ fn test_take_reverse(array: &dyn Array) { let len = array.len(); // Take elements in reverse order let indices = PrimitiveArray::from_iter((0..len as u64).rev()); - let result = - take(array, indices.as_ref()).vortex_expect("take should succeed in conformance test"); + let result = array + .take(indices.to_array()) + .vortex_expect("take should succeed in conformance test"); assert_eq!(result.len(), len); @@ -259,8 +266,9 @@ fn test_take_single_middle(array: &dyn Array) { let middle_idx = len / 2; let indices = PrimitiveArray::from_iter([middle_idx as u64]); - let result = - take(array, indices.as_ref()).vortex_expect("take should succeed in conformance test"); + let result = array + .take(indices.to_array()) + .vortex_expect("take should succeed in conformance test"); assert_eq!(result.len(), 1); assert_eq!( @@ -286,7 +294,8 @@ fn test_take_random_unsorted(array: &dyn Array) { } let indices_array = PrimitiveArray::from_iter(indices.clone()); - let result = take(array, indices_array.as_ref()) + let result = array + .take(indices_array.to_array()) .vortex_expect("take should succeed in conformance test"); assert_eq!(result.len(), indices.len()); @@ -311,8 +320,9 @@ fn test_take_contiguous_range(array: &dyn Array) { // Take a contiguous range from the middle let indices = PrimitiveArray::from_iter(start as u64..end as u64); - let result = - take(array, indices.as_ref()).vortex_expect("take should succeed in conformance test"); + let result = array + .take(indices.to_array()) + .vortex_expect("take should succeed in conformance test"); assert_eq!(result.len(), end - start); @@ -346,7 +356,8 @@ fn test_take_mixed_repeated(array: &dyn Array) { ]; let indices_array = PrimitiveArray::from_iter(indices.clone()); - let result = take(array, indices_array.as_ref()) + let result = array + .take(indices_array.to_array()) .vortex_expect("take should succeed in conformance test"); assert_eq!(result.len(), indices.len()); @@ -376,7 +387,8 @@ fn test_take_large_indices(array: &dyn Array) { .collect(); let indices_array = PrimitiveArray::from_iter(indices.clone()); - let result = take(array, indices_array.as_ref()) + let result = array + .take(indices_array.to_array()) .vortex_expect("take should succeed in conformance test"); assert_eq!(result.len(), num_indices); diff --git a/vortex-array/src/compute/mod.rs b/vortex-array/src/compute/mod.rs index b88b00413fa..7741947f17e 100644 --- a/vortex-array/src/compute/mod.rs +++ b/vortex-array/src/compute/mod.rs @@ -33,7 +33,6 @@ pub use nan_count::*; pub use numeric::*; use parking_lot::RwLock; pub use sum::*; -pub use take::*; use vortex_dtype::DType; use vortex_error::VortexError; use vortex_error::VortexResult; @@ -67,7 +66,6 @@ mod min_max; mod nan_count; mod numeric; mod sum; -mod take; mod zip; /// An instance of a compute function holding the implementation vtable and a set of registered diff --git a/vortex-array/src/compute/take.rs b/vortex-array/src/compute/take.rs deleted file mode 100644 index f0e49e9a764..00000000000 --- a/vortex-array/src/compute/take.rs +++ /dev/null @@ -1,17 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use vortex_error::VortexResult; - -use crate::Array; -use crate::ArrayRef; - -/// Creates a new array using the elements from the input `array` indexed by `indices`. -/// -/// For example, if we have an `array` `[1, 2, 3, 4, 5]` and `indices` `[4, 2]`, the resulting -/// array would be `[5, 3]`. -/// -/// The output array will have the same length as the `indices` array. -pub fn take(array: &dyn Array, indices: &dyn Array) -> VortexResult { - array.take(indices.to_array()) -} diff --git a/vortex-array/src/patches.rs b/vortex-array/src/patches.rs index 06d35912074..1df1efb8559 100644 --- a/vortex-array/src/patches.rs +++ b/vortex-array/src/patches.rs @@ -43,7 +43,6 @@ use crate::arrays::PrimitiveArray; use crate::compute::cast; use crate::compute::filter; use crate::compute::is_sorted; -use crate::compute::take; use crate::search_sorted::SearchResult; use crate::search_sorted::SearchSorted; use crate::search_sorted::SearchSortedSide; @@ -800,10 +799,9 @@ impl Patches { array_len: new_array_len, offset: 0, indices: new_indices.clone(), - values: take( - self.values(), - &PrimitiveArray::new(values_indices, values_validity).into_array(), - )?, + values: self + .values() + .take(PrimitiveArray::new(values_indices, values_validity).into_array())?, chunk_offsets: None, offset_within_chunk: Some(0), // Reset when creating new Patches. })) @@ -837,7 +835,7 @@ impl Patches { return Ok(None); }; - let taken_values = take(self.values(), &value_indices)?; + let taken_values = self.values().take(value_indices)?; Ok(Some(Patches { array_len: new_length,