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

18.18

/vortex-array/src/arrays/chunked/decode.rs

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

use vortex_buffer::BufferMut;
use vortex_dtype::{DType, Nullability, PType, StructFields};
use vortex_error::{VortexExpect, VortexResult, vortex_err};

use super::ChunkedArray;
use crate::arrays::{ChunkedVTable, ListArray, PrimitiveArray, StructArray};
use crate::builders::{ArrayBuilder, builder_with_capacity};
use crate::compute::cast;
use crate::validity::Validity;
use crate::vtable::CanonicalVTable;
use crate::{Array as _, ArrayRef, Canonical, IntoArray, ToCanonical};

impl CanonicalVTable<ChunkedVTable> for ChunkedVTable {
    fn canonicalize(array: &ChunkedArray) -> VortexResult<Canonical> {
        if array.nchunks() == 0 {
            return Ok(Canonical::empty(array.dtype()));
        }
        if array.nchunks() == 1 {
            return array.chunks()[0].to_canonical();
        }
        match array.dtype() {
            DType::Struct(struct_dtype, _) => {
                let struct_array = swizzle_struct_chunks(
                    array.chunks(),
                    Validity::copy_from_array(array.as_ref())?,
                    struct_dtype,
                )?;
                Ok(Canonical::Struct(struct_array))
            }
            DType::List(elem, _) => Ok(Canonical::List(pack_lists(
                array.chunks(),
                Validity::copy_from_array(array.as_ref())?,
                elem,
            )?)),
            _ => {
                let mut builder = builder_with_capacity(array.dtype(), array.len());
                array.append_to_builder(builder.as_mut())?;
                builder.finish().to_canonical()
            }
        }
    }

    fn append_to_builder(array: &ChunkedArray, builder: &mut dyn ArrayBuilder) -> VortexResult<()> {
        for chunk in array.chunks() {
            chunk.append_to_builder(builder)?;
        }
        Ok(())
    }
}

/// Swizzle the pointers within a ChunkedArray of StructArrays to instead be a single
/// StructArray, where the Array for each Field is a ChunkedArray.
fn swizzle_struct_chunks(
    chunks: &[ArrayRef],
    validity: Validity,
    struct_dtype: &StructFields,
) -> VortexResult<StructArray> {
    let len = chunks.iter().map(|chunk| chunk.len()).sum();
    let mut field_arrays = Vec::new();

    for (field_idx, field_dtype) in struct_dtype.fields().enumerate() {
        let field_chunks = chunks
            .iter()
            .map(|c| {
                c.to_struct()
                    .vortex_expect("Chunk was not a StructArray")
                    .fields()
                    .get(field_idx)
                    .vortex_expect("Invalid field index")
                    .to_array()
            })
            .collect::<Vec<_>>();
        let field_array = ChunkedArray::try_new(field_chunks, field_dtype.clone())?;
        field_arrays.push(field_array.into_array());
    }

    StructArray::try_new_with_dtype(field_arrays, struct_dtype.clone(), len, validity)
}

fn pack_lists(
    chunks: &[ArrayRef],
    validity: Validity,
    elem_dtype: &DType,
) -> VortexResult<ListArray> {
    let len: usize = chunks.iter().map(|c| c.len()).sum();
    let mut offsets = BufferMut::<i64>::with_capacity(len + 1);
    offsets.push(0);
    let mut elements = Vec::new();

    for chunk in chunks {
        let chunk = chunk.to_list()?;
        // TODO: handle i32 offsets if they fit.
        let offsets_arr = cast(
            chunk.offsets(),
            &DType::Primitive(PType::I64, Nullability::NonNullable),
        )?
        .to_primitive()?;

        let first_offset_value: usize = usize::try_from(&offsets_arr.scalar_at(0)?)?;
        let last_offset_value: usize =
            usize::try_from(&offsets_arr.scalar_at(offsets_arr.len() - 1)?)?;
        elements.push(
            chunk
                .elements()
                .slice(first_offset_value, last_offset_value)?,
        );

        let adjustment_from_previous = *offsets
            .last()
            .ok_or_else(|| vortex_err!("List offsets must have at least one element"))?;
        offsets.extend_trusted(
            offsets_arr
                .as_slice::<i64>()
                .iter()
                .skip(1)
                .map(|off| off + adjustment_from_previous - first_offset_value as i64),
        );
    }
    let chunked_elements = ChunkedArray::try_new(elements, elem_dtype.clone())?.into_array();
    let offsets = PrimitiveArray::new(offsets.freeze(), Validity::NonNullable);

    ListArray::try_new(chunked_elements, offsets.into_array(), validity)
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use vortex_dtype::DType::{List, Primitive};
    use vortex_dtype::Nullability::NonNullable;
    use vortex_dtype::PType::I32;

    use crate::accessor::ArrayAccessor;
    use crate::arrays::{ChunkedArray, ListArray, PrimitiveArray, StructArray, VarBinViewArray};
    use crate::validity::Validity;
    use crate::{IntoArray, ToCanonical};

    #[test]
    pub fn pack_nested_structs() {
        let struct_array = StructArray::try_new(
            vec!["a".into()].into(),
            vec![VarBinViewArray::from_iter_str(["foo", "bar", "baz", "quak"]).into_array()],
            4,
            Validity::NonNullable,
        )
        .unwrap();
        let dtype = struct_array.dtype().clone();
        let chunked = ChunkedArray::try_new(
            vec![
                ChunkedArray::try_new(vec![struct_array.to_array()], dtype.clone())
                    .unwrap()
                    .into_array(),
            ],
            dtype,
        )
        .unwrap()
        .into_array();
        let canonical_struct = chunked.to_struct().unwrap();
        let canonical_varbin = canonical_struct.fields()[0].to_varbinview().unwrap();
        let original_varbin = struct_array.fields()[0].to_varbinview().unwrap();
        let orig_values = original_varbin
            .with_iterator(|it| it.map(|a| a.map(|v| v.to_vec())).collect::<Vec<_>>())
            .unwrap();
        let canon_values = canonical_varbin
            .with_iterator(|it| it.map(|a| a.map(|v| v.to_vec())).collect::<Vec<_>>())
            .unwrap();
        assert_eq!(orig_values, canon_values);
    }

    #[test]
    pub fn pack_nested_lists() {
        let l1 = ListArray::try_new(
            PrimitiveArray::from_iter([1, 2, 3, 4]).into_array(),
            PrimitiveArray::from_iter([0, 3]).into_array(),
            Validity::NonNullable,
        )
        .unwrap();

        let l2 = ListArray::try_new(
            PrimitiveArray::from_iter([5, 6]).into_array(),
            PrimitiveArray::from_iter([0, 2]).into_array(),
            Validity::NonNullable,
        )
        .unwrap();

        let chunked_list = ChunkedArray::try_new(
            vec![l1.clone().into_array(), l2.clone().into_array()],
            List(Arc::new(Primitive(I32, NonNullable)), NonNullable),
        );

        let canon_values = chunked_list.unwrap().to_list().unwrap();

        assert_eq!(l1.scalar_at(0).unwrap(), canon_values.scalar_at(0).unwrap());
        assert_eq!(l2.scalar_at(0).unwrap(), canon_values.scalar_at(1).unwrap());
    }
}

1	// SPDX-License-Identifier: Apache-2.0
2	// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4	use vortex_buffer::BufferMut;
5	use vortex_dtype::{DType, Nullability, PType, StructFields};
6	use vortex_error::{VortexExpect, VortexResult, vortex_err};
7
8	use super::ChunkedArray;
9	use crate::arrays::{ChunkedVTable, ListArray, PrimitiveArray, StructArray};
10	use crate::builders::{ArrayBuilder, builder_with_capacity};
11	use crate::compute::cast;
12	use crate::validity::Validity;
13	use crate::vtable::CanonicalVTable;
14	use crate::{Array as _, ArrayRef, Canonical, IntoArray, ToCanonical};
15
16	impl CanonicalVTable<ChunkedVTable> for ChunkedVTable {
17	fn canonicalize(array: &ChunkedArray) -> VortexResult<Canonical> {	24✔
18	if array.nchunks() == 0 {	24✔
UNCOV 19	return Ok(Canonical::empty(array.dtype()));	×
20	}	24✔
21	if array.nchunks() == 1 {	24✔
22	return array.chunks()[0].to_canonical();	4✔
23	}	20✔
24	match array.dtype() {	20✔
25	DType::Struct(struct_dtype, _) => {	×
26	let struct_array = swizzle_struct_chunks(	×
27	array.chunks(),	×
28	Validity::copy_from_array(array.as_ref())?,	×
29	struct_dtype,	×
30	)?;	×
31	Ok(Canonical::Struct(struct_array))	×
32	}
UNCOV 33	DType::List(elem, _) => Ok(Canonical::List(pack_lists(	×
UNCOV 34	array.chunks(),	×
UNCOV 35	Validity::copy_from_array(array.as_ref())?,	×
UNCOV 36	elem,	×
37	)?)),	×
38	_ => {
39	let mut builder = builder_with_capacity(array.dtype(), array.len());	20✔
40	array.append_to_builder(builder.as_mut())?;	20✔
41	builder.finish().to_canonical()	20✔
42	}
43	}
44	}	24✔
45
46	fn append_to_builder(array: &ChunkedArray, builder: &mut dyn ArrayBuilder) -> VortexResult<()> {	20✔
47	for chunk in array.chunks() {	172✔
48	chunk.append_to_builder(builder)?;	172✔
49	}
50	Ok(())	20✔
51	}	20✔
52	}
53
54	/// Swizzle the pointers within a ChunkedArray of StructArrays to instead be a single
55	/// StructArray, where the Array for each Field is a ChunkedArray.
56	fn swizzle_struct_chunks(	×
57	chunks: &[ArrayRef],	×
58	validity: Validity,	×
59	struct_dtype: &StructFields,	×
60	) -> VortexResult<StructArray> {	×
61	let len = chunks.iter().map(\|chunk\| chunk.len()).sum();	×
62	let mut field_arrays = Vec::new();	×
63
64	for (field_idx, field_dtype) in struct_dtype.fields().enumerate() {	×
65	let field_chunks = chunks	×
66	.iter()	×
67	.map(\|c\| {	×
68	c.to_struct()	×
69	.vortex_expect("Chunk was not a StructArray")	×
70	.fields()	×
71	.get(field_idx)	×
72	.vortex_expect("Invalid field index")	×
73	.to_array()	×
74	})	×
75	.collect::<Vec<_>>();	×
76	let field_array = ChunkedArray::try_new(field_chunks, field_dtype.clone())?;	×
77	field_arrays.push(field_array.into_array());	×
78	}
79
80	StructArray::try_new_with_dtype(field_arrays, struct_dtype.clone(), len, validity)	×
81	}	×
82
UNCOV 83	fn pack_lists(	×
UNCOV 84	chunks: &[ArrayRef],	×
UNCOV 85	validity: Validity,	×
UNCOV 86	elem_dtype: &DType,	×
UNCOV 87	) -> VortexResult<ListArray> {	×
UNCOV 88	let len: usize = chunks.iter().map(\|c\| c.len()).sum();	×
UNCOV 89	let mut offsets = BufferMut::<i64>::with_capacity(len + 1);	×
UNCOV 90	offsets.push(0);	×
UNCOV 91	let mut elements = Vec::new();	×
92
UNCOV 93	for chunk in chunks {	×
UNCOV 94	let chunk = chunk.to_list()?;	×
95	// TODO: handle i32 offsets if they fit.
UNCOV 96	let offsets_arr = cast(	×
UNCOV 97	chunk.offsets(),	×
UNCOV 98	&DType::Primitive(PType::I64, Nullability::NonNullable),	×
99	)?	×
UNCOV 100	.to_primitive()?;	×
101
UNCOV 102	let first_offset_value: usize = usize::try_from(&offsets_arr.scalar_at(0)?)?;	×
UNCOV 103	let last_offset_value: usize =	×
UNCOV 104	usize::try_from(&offsets_arr.scalar_at(offsets_arr.len() - 1)?)?;	×
UNCOV 105	elements.push(	×
UNCOV 106	chunk	×
UNCOV 107	.elements()	×
UNCOV 108	.slice(first_offset_value, last_offset_value)?,	×
109	);
110
UNCOV 111	let adjustment_from_previous = *offsets	×
UNCOV 112	.last()	×
UNCOV 113	.ok_or_else(\|\| vortex_err!("List offsets must have at least one element"))?;	×
UNCOV 114	offsets.extend_trusted(	×
UNCOV 115	offsets_arr	×
UNCOV 116	.as_slice::<i64>()	×
UNCOV 117	.iter()	×
UNCOV 118	.skip(1)	×
UNCOV 119	.map(\|off\| off + adjustment_from_previous - first_offset_value as i64),	×
120	);
121	}
UNCOV 122	let chunked_elements = ChunkedArray::try_new(elements, elem_dtype.clone())?.into_array();	×
UNCOV 123	let offsets = PrimitiveArray::new(offsets.freeze(), Validity::NonNullable);	×
124
UNCOV 125	ListArray::try_new(chunked_elements, offsets.into_array(), validity)	×
UNCOV 126	}	×
127
128	#[cfg(test)]
129	mod tests {
130	use std::sync::Arc;
131
132	use vortex_dtype::DType::{List, Primitive};
133	use vortex_dtype::Nullability::NonNullable;
134	use vortex_dtype::PType::I32;
135
136	use crate::accessor::ArrayAccessor;
137	use crate::arrays::{ChunkedArray, ListArray, PrimitiveArray, StructArray, VarBinViewArray};
138	use crate::validity::Validity;
139	use crate::{IntoArray, ToCanonical};
140
141	#[test]
142	pub fn pack_nested_structs() {
143	let struct_array = StructArray::try_new(
144	vec!["a".into()].into(),
145	vec![VarBinViewArray::from_iter_str(["foo", "bar", "baz", "quak"]).into_array()],
146	4,
147	Validity::NonNullable,
148	)
149	.unwrap();
150	let dtype = struct_array.dtype().clone();
151	let chunked = ChunkedArray::try_new(
152	vec![
153	ChunkedArray::try_new(vec![struct_array.to_array()], dtype.clone())
154	.unwrap()
155	.into_array(),
156	],
157	dtype,
158	)
159	.unwrap()
160	.into_array();
161	let canonical_struct = chunked.to_struct().unwrap();
162	let canonical_varbin = canonical_struct.fields()[0].to_varbinview().unwrap();
163	let original_varbin = struct_array.fields()[0].to_varbinview().unwrap();
164	let orig_values = original_varbin
165	.with_iterator(\|it\| it.map(\|a\| a.map(\|v\| v.to_vec())).collect::<Vec<_>>())
166	.unwrap();
167	let canon_values = canonical_varbin
168	.with_iterator(\|it\| it.map(\|a\| a.map(\|v\| v.to_vec())).collect::<Vec<_>>())
169	.unwrap();
170	assert_eq!(orig_values, canon_values);
171	}
172
173	#[test]
174	pub fn pack_nested_lists() {
175	let l1 = ListArray::try_new(
176	PrimitiveArray::from_iter([1, 2, 3, 4]).into_array(),
177	PrimitiveArray::from_iter([0, 3]).into_array(),
178	Validity::NonNullable,
179	)
180	.unwrap();
181
182	let l2 = ListArray::try_new(
183	PrimitiveArray::from_iter([5, 6]).into_array(),
184	PrimitiveArray::from_iter([0, 2]).into_array(),
185	Validity::NonNullable,
186	)
187	.unwrap();
188
189	let chunked_list = ChunkedArray::try_new(
190	vec![l1.clone().into_array(), l2.clone().into_array()],
191	List(Arc::new(Primitive(I32, NonNullable)), NonNullable),
192	);
193
194	let canon_values = chunked_list.unwrap().to_list().unwrap();
195
196	assert_eq!(l1.scalar_at(0).unwrap(), canon_values.scalar_at(0).unwrap());
197	assert_eq!(l2.scalar_at(0).unwrap(), canon_values.scalar_at(1).unwrap());
198	}
199	}

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