17051909105

Committed 18 Aug 2025 08:37PM UTC coverage: 87.884% (+0.02%) from 87.865%

Build # 17051909105

Build Type

Pull #4278

github

Committed by

web-flow

Commit Message

Merge ab4e1c32c into 18bd3ddc0

Pull Request Pull Request #4278: Chore: Avoid a clone in `take`

Run Details

7 of 7 new or added lines in 1 file covered. (100.0%)

3 existing lines in 1 file now uncovered.

56584 of 64385 relevant lines covered (87.88%)

628649.98 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

87.25

/vortex-array/src/compute/take.rs

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

use std::sync::LazyLock;

use arcref::ArcRef;
use vortex_dtype::DType;
use vortex_error::{VortexError, VortexResult, vortex_bail, vortex_err};
use vortex_scalar::Scalar;

use crate::arrays::ConstantArray;
use crate::compute::{ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Output};
use crate::stats::{Precision, Stat, StatsSet};
use crate::vtable::VTable;
use crate::{Array, ArrayRef, Canonical, IntoArray};

static TAKE_FN: LazyLock<ComputeFn> = LazyLock::new(|| {
    let compute = ComputeFn::new("take".into(), ArcRef::new_ref(&Take));
    for kernel in inventory::iter::<TakeKernelRef> {
        compute.register_kernel(kernel.0.clone());
    }
    compute
});

/// 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<ArrayRef> {
    if indices.is_empty() {
        return Ok(Canonical::empty(
            &array
                .dtype()
                .union_nullability(indices.dtype().nullability()),
        )
        .into_array());
    }

    TAKE_FN
        .invoke(&InvocationArgs {
            inputs: &[array.into(), indices.into()],
            options: &(),
        })?
        .unwrap_array()
}

#[doc(hidden)]
pub struct Take;

impl ComputeFnVTable for Take {
    fn invoke(
        &self,
        args: &InvocationArgs,
        kernels: &[ArcRef<dyn Kernel>],
    ) -> VortexResult<Output> {
        let TakeArgs { array, indices } = TakeArgs::try_from(args)?;

        // TODO(ngates): if indices are sorted and unique (strict-sorted), then we should delegate to
        //  the filter function since they're typically optimised for this case.
        // TODO(ngates): if indices min is quite high, we could slice self and offset the indices
        //  such that canonicalize does less work.

        if indices.all_invalid()? {
            return Ok(ConstantArray::new(
                Scalar::null(array.dtype().as_nullable()),
                indices.len(),
            )
            .into_array()
            .into());
        }

        // We know that constant array don't need stats propagation, so we can avoid the overhead of
        // computing derived stats and merging them in.
        let derived_stats = (!array.is_constant()).then(|| derive_take_stats(array));

        let taken = take_impl(array, indices, kernels)?;

        if let Some(derived_stats) = derived_stats {
            let mut stats = taken.statistics().to_owned();
            stats.combine_sets(&derived_stats, array.dtype())?;
            for (stat, val) in stats.into_iter() {
                taken.statistics().set(stat, val)
            }
        }

        Ok(taken.into())
    }

    fn return_dtype(&self, args: &InvocationArgs) -> VortexResult<DType> {
        let TakeArgs { array, indices } = TakeArgs::try_from(args)?;

        if !indices.dtype().is_int() {
            vortex_bail!(
                "Take indices must be an integer type, got {}",
                indices.dtype()
            );
        }

        Ok(array
            .dtype()
            .union_nullability(indices.dtype().nullability()))
    }

    fn return_len(&self, args: &InvocationArgs) -> VortexResult<usize> {
        let TakeArgs { indices, .. } = TakeArgs::try_from(args)?;
        Ok(indices.len())
    }

    fn is_elementwise(&self) -> bool {
        false
    }
}

fn derive_take_stats(arr: &dyn Array) -> StatsSet {
    let stats = arr.statistics().to_owned();

    let is_constant = arr.statistics().get_as::<bool>(Stat::IsConstant);

    let mut stats = stats.keep_inexact_stats(&[
        // Cannot create values smaller than min or larger than max
        Stat::Min,
        Stat::Max,
    ]);

    if is_constant == Some(Precision::Exact(true)) {
        // Any combination of elements from a constant array is still const
        stats.set(Stat::IsConstant, Precision::exact(true));
    }

    stats
}

fn take_impl(
    array: &dyn Array,
    indices: &dyn Array,
    kernels: &[ArcRef<dyn Kernel>],
) -> VortexResult<ArrayRef> {
    let args = InvocationArgs {
        inputs: &[array.into(), indices.into()],
        options: &(),
    };

    // First look for a TakeFrom specialized on the indices.
    for kernel in TAKE_FROM_FN.kernels() {
        if let Some(output) = kernel.invoke(&args)? {
            return output.unwrap_array();
        }
    }
    if let Some(output) = indices.invoke(&TAKE_FROM_FN, &args)? {
        return output.unwrap_array();
    }

    // Then look for a Take kernel
    for kernel in kernels {
        if let Some(output) = kernel.invoke(&args)? {
            return output.unwrap_array();
        }
    }
    if let Some(output) = array.invoke(&TAKE_FN, &args)? {
        return output.unwrap_array();
    }

    // Otherwise, canonicalize and try again.
    if !array.is_canonical() {
        log::debug!("No take implementation found for {}", array.encoding_id());
        let canonical = array.to_canonical()?;
        return take(canonical.as_ref(), indices);
    }

    vortex_bail!("No take implementation found for {}", array.encoding_id());
}

struct TakeArgs<'a> {
    array: &'a dyn Array,
    indices: &'a dyn Array,
}

impl<'a> TryFrom<&InvocationArgs<'a>> for TakeArgs<'a> {
    type Error = VortexError;

    fn try_from(value: &InvocationArgs<'a>) -> Result<Self, Self::Error> {
        if value.inputs.len() != 2 {
            vortex_bail!("Expected 2 inputs, found {}", value.inputs.len());
        }
        let array = value.inputs[0]
            .array()
            .ok_or_else(|| vortex_err!("Expected first input to be an array"))?;
        let indices = value.inputs[1]
            .array()
            .ok_or_else(|| vortex_err!("Expected second input to be an array"))?;
        Ok(Self { array, indices })
    }
}

pub trait TakeKernel: VTable {
    /// Create a new array by taking the values from the `array` at the
    /// given `indices`.
    ///
    /// # Panics
    ///
    /// Using `indices` that are invalid for the given `array` will cause a panic.
    fn take(&self, array: &Self::Array, indices: &dyn Array) -> VortexResult<ArrayRef>;
}

/// A kernel that implements the filter function.
pub struct TakeKernelRef(pub ArcRef<dyn Kernel>);
inventory::collect!(TakeKernelRef);

#[derive(Debug)]
pub struct TakeKernelAdapter<V: VTable>(pub V);

impl<V: VTable + TakeKernel> TakeKernelAdapter<V> {
    pub const fn lift(&'static self) -> TakeKernelRef {
        TakeKernelRef(ArcRef::new_ref(self))
    }
}

impl<V: VTable + TakeKernel> Kernel for TakeKernelAdapter<V> {
    fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {
        let inputs = TakeArgs::try_from(args)?;
        let Some(array) = inputs.array.as_opt::<V>() else {
            return Ok(None);
        };
        Ok(Some(V::take(&self.0, array, inputs.indices)?.into()))
    }
}

static TAKE_FROM_FN: LazyLock<ComputeFn> = LazyLock::new(|| {
    let compute = ComputeFn::new("take_from".into(), ArcRef::new_ref(&TakeFrom));
    for kernel in inventory::iter::<TakeFromKernelRef> {
        compute.register_kernel(kernel.0.clone());
    }
    compute
});

pub struct TakeFrom;

impl ComputeFnVTable for TakeFrom {
    fn invoke(
        &self,
        _args: &InvocationArgs,
        _kernels: &[ArcRef<dyn Kernel>],
    ) -> VortexResult<Output> {
        vortex_bail!(
            "TakeFrom should not be invoked directly. Its kernels are used to accelerated the Take function"
        )
    }

    fn return_dtype(&self, args: &InvocationArgs) -> VortexResult<DType> {
        Take.return_dtype(args)
    }

    fn return_len(&self, args: &InvocationArgs) -> VortexResult<usize> {
        Take.return_len(args)
    }

    fn is_elementwise(&self) -> bool {
        Take.is_elementwise()
    }
}

pub trait TakeFromKernel: VTable {
    /// Create a new array by taking the values from the `array` at the
    /// given `indices`.
    fn take_from(&self, indices: &Self::Array, array: &dyn Array)
    -> VortexResult<Option<ArrayRef>>;
}

pub struct TakeFromKernelRef(pub ArcRef<dyn Kernel>);
inventory::collect!(TakeFromKernelRef);

#[derive(Debug)]
pub struct TakeFromKernelAdapter<V: VTable>(pub V);

impl<V: VTable + TakeFromKernel> TakeFromKernelAdapter<V> {
    pub const fn lift(&'static self) -> TakeFromKernelRef {
        TakeFromKernelRef(ArcRef::new_ref(self))
    }
}

impl<V: VTable + TakeFromKernel> Kernel for TakeFromKernelAdapter<V> {
    fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {
        let inputs = TakeArgs::try_from(args)?;
        let Some(indices) = inputs.indices.as_opt::<V>() else {
            return Ok(None);
        };
        Ok(V::take_from(&self.0, indices, inputs.array)?.map(Output::from))
    }
}

1	// SPDX-License-Identifier: Apache-2.0
2	// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4	use std::sync::LazyLock;
5
6	use arcref::ArcRef;
7	use vortex_dtype::DType;
8	use vortex_error::{VortexError, VortexResult, vortex_bail, vortex_err};
9	use vortex_scalar::Scalar;
10
11	use crate::arrays::ConstantArray;
12	use crate::compute::{ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Output};
13	use crate::stats::{Precision, Stat, StatsSet};
14	use crate::vtable::VTable;
15	use crate::{Array, ArrayRef, Canonical, IntoArray};
16
17	static TAKE_FN: LazyLock<ComputeFn> = LazyLock::new(\|\| {	9,166✔
18	let compute = ComputeFn::new("take".into(), ArcRef::new_ref(&Take));	9,166✔
19	for kernel in inventory::iter::<TakeKernelRef> {	130,503✔
20	compute.register_kernel(kernel.0.clone());	121,337✔
21	}	121,337✔
22	compute	9,166✔
23	});	9,166✔
24
25	/// Creates a new array using the elements from the input `array` indexed by `indices`.
26	///
27	/// For example, if we have an `array` `[1, 2, 3, 4, 5]` and `indices` `[4, 2]`, the resulting
28	/// array would be `[5, 3]`.
29	///
30	/// The output array will have the same length as the `indices` array.
31	pub fn take(array: &dyn Array, indices: &dyn Array) -> VortexResult<ArrayRef> {	129,434✔
32	if indices.is_empty() {	129,434✔
33	return Ok(Canonical::empty(	9,122✔
34	&array	9,122✔
35	.dtype()	9,122✔
36	.union_nullability(indices.dtype().nullability()),	9,122✔
37	)	9,122✔
38	.into_array());	9,122✔
39	}	120,312✔
40
41	TAKE_FN	120,312✔
42	.invoke(&InvocationArgs {	120,312✔
43	inputs: &[array.into(), indices.into()],	120,312✔
44	options: &(),	120,312✔
45	})?	120,312✔
46	.unwrap_array()	120,273✔
47	}	129,434✔
48
49	#[doc(hidden)]
50	pub struct Take;
51
52	impl ComputeFnVTable for Take {
53	fn invoke(	120,312✔
54	&self,	120,312✔
55	args: &InvocationArgs,	120,312✔
56	kernels: &[ArcRef<dyn Kernel>],	120,312✔
57	) -> VortexResult<Output> {	120,312✔
58	let TakeArgs { array, indices } = TakeArgs::try_from(args)?;	120,312✔
59
60	// TODO(ngates): if indices are sorted and unique (strict-sorted), then we should delegate to
61	// the filter function since they're typically optimised for this case.
62	// TODO(ngates): if indices min is quite high, we could slice self and offset the indices
63	// such that canonicalize does less work.
64
65	if indices.all_invalid()? {	120,312✔
66	return Ok(ConstantArray::new(	481✔
67	Scalar::null(array.dtype().as_nullable()),	481✔
68	indices.len(),	481✔
69	)	481✔
70	.into_array()	481✔
71	.into());	481✔
72	}	119,831✔
73
74	// We know that constant array don't need stats propagation, so we can avoid the overhead of
75	// computing derived stats and merging them in.
76	let derived_stats = (!array.is_constant()).then(\|\| derive_take_stats(array));	119,831✔
77
78	let taken = take_impl(array, indices, kernels)?;	119,831✔
79
80	if let Some(derived_stats) = derived_stats {	119,792✔
81	let mut stats = taken.statistics().to_owned();	101,915✔
82	stats.combine_sets(&derived_stats, array.dtype())?;	101,915✔
83	for (stat, val) in stats.into_iter() {	102,215✔
84	taken.statistics().set(stat, val)	2,578✔
85	}
86	}	17,877✔
87
88	Ok(taken.into())	119,792✔
89	}	120,312✔
90
91	fn return_dtype(&self, args: &InvocationArgs) -> VortexResult<DType> {	120,312✔
92	let TakeArgs { array, indices } = TakeArgs::try_from(args)?;	120,312✔
93
94	if !indices.dtype().is_int() {	120,312✔
95	vortex_bail!(	×
96	"Take indices must be an integer type, got {}",
UNCOV 97	indices.dtype()	×
98	);
99	}	120,312✔
100
101	Ok(array	120,312✔
102	.dtype()	120,312✔
103	.union_nullability(indices.dtype().nullability()))	120,312✔
104	}	120,312✔
105
106	fn return_len(&self, args: &InvocationArgs) -> VortexResult<usize> {	120,312✔
107	let TakeArgs { indices, .. } = TakeArgs::try_from(args)?;	120,312✔
108	Ok(indices.len())	120,312✔
109	}	120,312✔
110
111	fn is_elementwise(&self) -> bool {	120,312✔
112	false	120,312✔
113	}	120,312✔
114	}
115
116	fn derive_take_stats(arr: &dyn Array) -> StatsSet {	101,954✔
117	let stats = arr.statistics().to_owned();	101,954✔
118
119	let is_constant = arr.statistics().get_as::<bool>(Stat::IsConstant);	101,954✔
120
121	let mut stats = stats.keep_inexact_stats(&[	101,954✔
122	// Cannot create values smaller than min or larger than max	101,954✔
123	Stat::Min,	101,954✔
124	Stat::Max,	101,954✔
125	]);	101,954✔
126
127	if is_constant == Some(Precision::Exact(true)) {	101,954✔
UNCOV 128	// Any combination of elements from a constant array is still const	×
129	stats.set(Stat::IsConstant, Precision::exact(true));
130	}	101,954✔
131
132	stats	101,954✔
133	}	101,954✔
134
135	fn take_impl(	119,831✔
136	array: &dyn Array,	119,831✔
137	indices: &dyn Array,	119,831✔
138	kernels: &[ArcRef<dyn Kernel>],	119,831✔
139	) -> VortexResult<ArrayRef> {	119,831✔
140	let args = InvocationArgs {	119,831✔
141	inputs: &[array.into(), indices.into()],	119,831✔
142	options: &(),	119,831✔
143	};	119,831✔
144
145	// First look for a TakeFrom specialized on the indices.
146	for kernel in TAKE_FROM_FN.kernels() {	119,831✔
147	if let Some(output) = kernel.invoke(&args)? {	14,037✔
148	return output.unwrap_array();	24✔
149	}	14,013✔
150	}
151	if let Some(output) = indices.invoke(&TAKE_FROM_FN, &args)? {	119,807✔
152	return output.unwrap_array();
153	}	119,807✔
154
155	// Then look for a Take kernel
156	for kernel in kernels {	1,035,518✔
157	if let Some(output) = kernel.invoke(&args)? {	1,029,746✔
158	return output.unwrap_array();	113,996✔
159	}	915,711✔
160	}
161	if let Some(output) = array.invoke(&TAKE_FN, &args)? {	5,772✔
162	return output.unwrap_array();
163	}	5,772✔
164
165	// Otherwise, canonicalize and try again.
166	if !array.is_canonical() {	5,772✔
167	log::debug!("No take implementation found for {}", array.encoding_id());	5,772✔
168	let canonical = array.to_canonical()?;	5,772✔
169	return take(canonical.as_ref(), indices);	5,772✔
170	}	×
171
172	vortex_bail!("No take implementation found for {}", array.encoding_id());
173	}	119,831✔
174
175	struct TakeArgs<'a> {
176	array: &'a dyn Array,
177	indices: &'a dyn Array,
178	}
179
180	impl<'a> TryFrom<&InvocationArgs<'a>> for TakeArgs<'a> {
181	type Error = VortexError;
182
183	fn try_from(value: &InvocationArgs<'a>) -> Result<Self, Self::Error> {	1,404,719✔
184	if value.inputs.len() != 2 {	1,404,719✔
UNCOV 185	vortex_bail!("Expected 2 inputs, found {}", value.inputs.len());	×
186	}	1,404,719✔
187	let array = value.inputs[0]	1,404,719✔
188	.array()	1,404,719✔
189	.ok_or_else(\|\| vortex_err!("Expected first input to be an array"))?;	1,404,719✔
190	let indices = value.inputs[1]	1,404,719✔
191	.array()	1,404,719✔
192	.ok_or_else(\|\| vortex_err!("Expected second input to be an array"))?;	1,404,719✔
193	Ok(Self { array, indices })	1,404,719✔
194	}	1,404,719✔
195	}
196
197	pub trait TakeKernel: VTable {
198	/// Create a new array by taking the values from the `array` at the
199	/// given `indices`.
200	///
201	/// # Panics
202	///
203	/// Using `indices` that are invalid for the given `array` will cause a panic.
204	fn take(&self, array: &Self::Array, indices: &dyn Array) -> VortexResult<ArrayRef>;
205	}
206
207	/// A kernel that implements the filter function.
208	pub struct TakeKernelRef(pub ArcRef<dyn Kernel>);
209	inventory::collect!(TakeKernelRef);
210
211	#[derive(Debug)]
212	pub struct TakeKernelAdapter<V: VTable>(pub V);
213
214	impl<V: VTable + TakeKernel> TakeKernelAdapter<V> {
215	pub const fn lift(&'static self) -> TakeKernelRef {	×
216	TakeKernelRef(ArcRef::new_ref(self))
217	}
218	}
219
220	impl<V: VTable + TakeKernel> Kernel for TakeKernelAdapter<V> {
221	fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {	965,167✔
222	let inputs = TakeArgs::try_from(args)?;	965,167✔
223	let Some(array) = inputs.array.as_opt::<V>() else {	965,167✔
224	return Ok(None);	890,442✔
225	};
226	Ok(Some(V::take(&self.0, array, inputs.indices)?.into()))	74,725✔
227	}	965,167✔
228	}
229
230	static TAKE_FROM_FN: LazyLock<ComputeFn> = LazyLock::new(\|\| {	9,125✔
231	let compute = ComputeFn::new("take_from".into(), ArcRef::new_ref(&TakeFrom));	9,125✔
232	for kernel in inventory::iter::<TakeFromKernelRef> {	10,415✔
233	compute.register_kernel(kernel.0.clone());	1,290✔
234	}	1,290✔
235	compute	9,125✔
236	});	9,125✔
237
238	pub struct TakeFrom;
239
240	impl ComputeFnVTable for TakeFrom {
241	fn invoke(	×
242	&self,	×
243	_args: &InvocationArgs,	×
244	_kernels: &[ArcRef<dyn Kernel>],	×
245	) -> VortexResult<Output> {	×
246	vortex_bail!(
247	"TakeFrom should not be invoked directly. Its kernels are used to accelerated the Take function"	×
248	)
249	}	×
250
251	fn return_dtype(&self, args: &InvocationArgs) -> VortexResult<DType> {	×
252	Take.return_dtype(args)
253	}	×
254
255	fn return_len(&self, args: &InvocationArgs) -> VortexResult<usize> {	×
256	Take.return_len(args)
257	}	×
258
259	fn is_elementwise(&self) -> bool {	×
260	Take.is_elementwise()
261	}
262	}
263
264	pub trait TakeFromKernel: VTable {
265	/// Create a new array by taking the values from the `array` at the
266	/// given `indices`.
267	fn take_from(&self, indices: &Self::Array, array: &dyn Array)
268	-> VortexResult<Option<ArrayRef>>;
269	}
270
271	pub struct TakeFromKernelRef(pub ArcRef<dyn Kernel>);
272	inventory::collect!(TakeFromKernelRef);
273
274	#[derive(Debug)]
275	pub struct TakeFromKernelAdapter<V: VTable>(pub V);
276
277	impl<V: VTable + TakeFromKernel> TakeFromKernelAdapter<V> {
278	pub const fn lift(&'static self) -> TakeFromKernelRef {	×
279	TakeFromKernelRef(ArcRef::new_ref(self))
280	}
281	}
282
283	impl<V: VTable + TakeFromKernel> Kernel for TakeFromKernelAdapter<V> {
284	fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {	3,135✔
285	let inputs = TakeArgs::try_from(args)?;	3,135✔
286	let Some(indices) = inputs.indices.as_opt::<V>() else {	3,135✔
287	return Ok(None);	3,111✔
288	};
289	Ok(V::take_from(&self.0, indices, inputs.array)?.map(Output::from))	24✔
290	}	3,135✔
291	}

vortex-data / vortex / 17051909105

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous