vortex_array/arrays/variant/vtable/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4mod kernel;
5mod operations;
6mod validity;
7
8use prost::Message;
9use vortex_error::VortexExpect;
10use vortex_error::VortexResult;
11use vortex_error::vortex_ensure;
12use vortex_error::vortex_panic;
13use vortex_proto::dtype as pb;
14use vortex_session::VortexSession;
15use vortex_session::registry::CachedId;
16use vortex_utils::aliases::hash_set::HashSet;
17
18use crate::ArrayRef;
19use crate::ExecutionCtx;
20use crate::ExecutionResult;
21use crate::array::Array;
22use crate::array::ArrayId;
23use crate::array::ArrayParts;
24use crate::array::ArrayView;
25use crate::array::EmptyArrayData;
26use crate::array::VTable;
27use crate::array::with_empty_buffers;
28use crate::arrays::variant::VariantSlots;
29use crate::arrays::variant::compute::rules::RULES;
30use crate::buffer::BufferHandle;
31use crate::dtype::DType;
32use crate::dtype::FieldName;
33use crate::dtype::FieldNames;
34use crate::dtype::Nullability;
35use crate::dtype::StructFields;
36use crate::scalar::Scalar;
37use crate::scalar::ScalarValue;
38use crate::serde::ArrayChildren;
39
40/// A [`Variant`]-encoded Vortex array.
41pub type VariantArray = Array<Variant>;
42
43pub(crate) fn initialize(session: &VortexSession) {
44    kernel::initialize(session);
45}
46
47#[derive(Clone, Debug)]
48pub struct Variant;
49
50#[derive(Clone, prost::Message)]
51struct VariantMetadataProto {
52    #[prost(message, optional, tag = "1")]
53    pub shredded_dtype: Option<pb::DType>,
54}
55
56impl VTable for Variant {
57    type TypedArrayData = EmptyArrayData;
58
59    type OperationsVTable = Self;
60
61    type ValidityVTable = Self;
62
63    fn id(&self) -> ArrayId {
64        static ID: CachedId = CachedId::new("vortex.variant");
65        *ID
66    }
67
68    fn validate(
69        &self,
70        _data: &Self::TypedArrayData,
71        dtype: &DType,
72        len: usize,
73        slots: &[Option<ArrayRef>],
74    ) -> VortexResult<()> {
75        vortex_ensure!(
76            slots.len() == VariantSlots::COUNT,
77            "VariantArray expects {} slots, got {}",
78            VariantSlots::COUNT,
79            slots.len()
80        );
81        vortex_ensure!(
82            slots[VariantSlots::CORE_STORAGE].is_some(),
83            "VariantArray core_storage slot must be present"
84        );
85        let core_storage = slots[VariantSlots::CORE_STORAGE]
86            .as_ref()
87            .vortex_expect("validated core_storage slot presence");
88        vortex_ensure!(
89            matches!(dtype, DType::Variant(_)),
90            "Expected Variant DType, got {dtype}"
91        );
92        vortex_ensure!(
93            core_storage.dtype() == dtype,
94            "VariantArray core_storage dtype {} does not match outer dtype {}",
95            core_storage.dtype(),
96            dtype
97        );
98        vortex_ensure!(
99            core_storage.len() == len,
100            "VariantArray core_storage length {} does not match outer length {}",
101            core_storage.len(),
102            len
103        );
104        if let Some(shredded) = slots[VariantSlots::SHREDDED].as_ref() {
105            vortex_ensure!(
106                shredded.len() == len,
107                "VariantArray shredded length {} does not match outer length {}",
108                shredded.len(),
109                len
110            );
111        }
112        Ok(())
113    }
114
115    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
116        0
117    }
118
119    fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
120        vortex_panic!("VariantArray buffer index {idx} out of bounds")
121    }
122
123    fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option<String> {
124        None
125    }
126
127    fn with_buffers(
128        &self,
129        array: ArrayView<'_, Self>,
130        buffers: &[BufferHandle],
131    ) -> VortexResult<ArrayParts<Self>> {
132        with_empty_buffers(self, array, buffers)
133    }
134
135    fn serialize(
136        array: ArrayView<'_, Self>,
137        _session: &VortexSession,
138    ) -> VortexResult<Option<Vec<u8>>> {
139        let shredded_dtype = array.slots()[VariantSlots::SHREDDED]
140            .as_ref()
141            .map(|shredded| shredded.dtype().try_into())
142            .transpose()?;
143        Ok(Some(
144            VariantMetadataProto { shredded_dtype }.encode_to_vec(),
145        ))
146    }
147
148    fn deserialize(
149        &self,
150        dtype: &DType,
151        len: usize,
152        metadata: &[u8],
153        buffers: &[BufferHandle],
154        children: &dyn ArrayChildren,
155        session: &VortexSession,
156    ) -> VortexResult<ArrayParts<Self>> {
157        vortex_ensure!(
158            buffers.is_empty(),
159            "VariantArray expects 0 buffers, got {}",
160            buffers.len()
161        );
162        let proto = VariantMetadataProto::decode(metadata)?;
163        let shredded_dtype = proto
164            .shredded_dtype
165            .as_ref()
166            .map(|dtype| DType::from_proto(dtype, session))
167            .transpose()?;
168        vortex_ensure!(matches!(dtype, DType::Variant(_)), "Expected Variant DType");
169        let expected_children = 1 + usize::from(shredded_dtype.is_some());
170        vortex_ensure!(
171            children.len() == expected_children,
172            "Expected {} children, got {}",
173            expected_children,
174            children.len(),
175        );
176        let core_storage = children.get(0, dtype, len)?;
177        let shredded = shredded_dtype
178            .map(|dtype| children.get(1, &dtype, len))
179            .transpose()?;
180        Ok(
181            ArrayParts::new(self.clone(), dtype.clone(), len, EmptyArrayData).with_slots(
182                VariantSlots {
183                    core_storage,
184                    shredded,
185                }
186                .into_slots(),
187            ),
188        )
189    }
190
191    fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
192        match VariantSlots::NAMES.get(idx) {
193            Some(name) => (*name).to_string(),
194            None => vortex_panic!("VariantArray slot_name index {idx} out of bounds"),
195        }
196    }
197
198    fn execute(array: Array<Self>, _ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
199        Ok(ExecutionResult::done(array))
200    }
201
202    fn reduce_parent(
203        array: ArrayView<'_, Self>,
204        parent: &ArrayRef,
205        child_idx: usize,
206    ) -> VortexResult<Option<ArrayRef>> {
207        RULES.evaluate(array, parent, child_idx)
208    }
209}
210
211fn merge_typed_scalar_as_variant(
212    typed_scalar: Scalar,
213    fallback_scalar: Option<Scalar>,
214    dtype: &DType,
215) -> VortexResult<Scalar> {
216    let scalar = if typed_scalar.is_null() {
217        fallback_scalar.unwrap_or_else(|| Scalar::null(dtype.clone()))
218    } else if matches!(
219        typed_scalar.dtype(),
220        DType::List(..) | DType::FixedSizeList(..)
221    ) {
222        Scalar::variant(typed_list_as_variant_payload(typed_scalar)?)
223    } else if typed_scalar.dtype().is_struct() {
224        merge_typed_object_as_variant(typed_scalar, fallback_scalar)?
225    } else if typed_scalar.dtype().is_variant() {
226        typed_scalar
227    } else {
228        Scalar::variant(typed_scalar)
229    };
230
231    if scalar.dtype() == dtype {
232        Ok(scalar)
233    } else {
234        scalar.cast(dtype)
235    }
236}
237
238fn typed_list_as_variant_payload(typed_scalar: Scalar) -> VortexResult<Scalar> {
239    let list = typed_scalar.as_list();
240    let elements = list
241        .elements()
242        .unwrap_or_default()
243        .into_iter()
244        .map(|element| {
245            if element.dtype().is_variant() {
246                element
247            } else {
248                Scalar::variant(element)
249            }
250        })
251        .collect();
252    Ok(Scalar::list(
253        DType::Variant(Nullability::NonNullable),
254        elements,
255        Nullability::NonNullable,
256    ))
257}
258
259fn merge_typed_object_as_variant(
260    typed_scalar: Scalar,
261    fallback_scalar: Option<Scalar>,
262) -> VortexResult<Scalar> {
263    let fallback_inner = fallback_scalar
264        .as_ref()
265        .and_then(|scalar| scalar.as_variant().value())
266        .filter(|scalar| scalar.dtype().is_struct() && !scalar.is_null());
267    let Some(fallback_inner) = fallback_inner else {
268        return Ok(Scalar::variant(typed_scalar));
269    };
270
271    merge_struct_payload(&typed_scalar, Some(fallback_inner)).map(Scalar::variant)
272}
273
274fn merge_struct_payload(typed: &Scalar, raw: Option<&Scalar>) -> VortexResult<Scalar> {
275    let typed_struct = typed.as_struct();
276    let raw_struct = raw
277        .filter(|scalar| scalar.dtype().is_struct() && !scalar.is_null())
278        .map(Scalar::as_struct);
279    let mut present_typed_fields = HashSet::new();
280    let mut names = Vec::new();
281    let mut values = Vec::new();
282
283    for name in typed_struct.names().iter() {
284        let Some(typed_field) = typed_struct.field(name.as_ref()) else {
285            continue;
286        };
287        if typed_field.is_null() {
288            continue;
289        }
290
291        let raw_field = raw_struct.and_then(|raw_struct| raw_struct.field(name.as_ref()));
292        let raw_payload = raw_field.as_ref().and_then(|scalar| {
293            if scalar.dtype().is_variant() {
294                scalar.as_variant().value()
295            } else {
296                Some(scalar)
297            }
298        });
299        let field = if typed_field.dtype().is_struct()
300            && raw_payload.is_some_and(|raw| raw.dtype().is_struct() && !raw.is_null())
301        {
302            Scalar::variant(merge_struct_payload(&typed_field, raw_payload)?)
303        } else if typed_field.dtype().is_variant() {
304            typed_field.cast(&DType::Variant(Nullability::NonNullable))?
305        } else {
306            Scalar::variant(typed_field)
307        };
308
309        present_typed_fields.insert(name.as_ref().to_string());
310        names.push(FieldName::from(name.as_ref()));
311        values.push(field.into_value());
312    }
313
314    if let Some(raw_struct) = raw_struct {
315        for name in raw_struct.names().iter() {
316            if present_typed_fields.contains(name.as_ref()) {
317                continue;
318            }
319            let Some(raw_field) = raw_struct.field(name.as_ref()) else {
320                continue;
321            };
322            if raw_field.is_null() {
323                continue;
324            }
325            let raw_field = if raw_field.dtype().is_variant() {
326                raw_field.cast(&DType::Variant(Nullability::NonNullable))?
327            } else {
328                Scalar::variant(raw_field)
329            };
330            names.push(FieldName::from(name.as_ref()));
331            values.push(raw_field.into_value());
332        }
333    }
334
335    let fields = StructFields::new(
336        FieldNames::from(names),
337        vec![DType::Variant(Nullability::NonNullable); values.len()],
338    );
339    Scalar::try_new(
340        DType::Struct(fields, Nullability::NonNullable),
341        Some(ScalarValue::Tuple(values)),
342    )
343}
344
345#[cfg(test)]
346mod tests {}
vortex_array/arrays/variant/vtable/mod.rs

vortex_array/arrays/variant/vtable/
mod.rs
