16935267080

Committed 13 Aug 2025 11:00AM UTC coverage: 24.312% (-63.3%) from 87.658%

Build # 16935267080

Build Type

Pull #4226

github

Committed by

web-flow

Commit Message

Merge 81b48c7fb into baa6ea202

Pull Request Pull Request #4226: Support converting TimestampTZ to and from duckdb

Run Details

0 of 2 new or added lines in 1 file covered. (0.0%)

20666 existing lines in 469 files now uncovered.

8726 of 35892 relevant lines covered (24.31%)

147.74 hits per line

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

0.0

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

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

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

vortex-data / vortex / 16935267080

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