chore[array]: inline take compute function

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,

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<u32> {
@@ -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())
 }

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)])

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);
     }
 

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)
         );
     }

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());
 }

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,26 +109,23 @@ 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);
     }
 
     #[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::<Canonical>(&mut LEGACY_SESSION.create_execution_ctx())
             .unwrap();
@@ -138,19 +134,17 @@ 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);
     }
 
     #[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());

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::<Canonical>(&mut LEGACY_SESSION.create_execution_ctx())
             .unwrap();