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

vortex_array/arrays/variant/vtable/
mod.rs
