vortex_arrow/
uuid.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Arrow plugin impls for the UUID extension type.
5//!
6//! UUIDs are a canonical Arrow extension type backed by `FixedSizeBinary[16]`. The Vortex side
7//! stores them as `FixedSizeList<u8; 16>`, so the conversion is a zero-copy reinterpretation
8//! of the byte buffer in both directions.
9
10use std::sync::Arc;
11
12use arrow_array::Array;
13use arrow_array::ArrayRef as ArrowArrayRef;
14use arrow_array::FixedSizeBinaryArray;
15use arrow_array::cast::AsArray;
16use arrow_array::types::UInt8Type;
17use arrow_schema::DataType;
18use arrow_schema::Field;
19use arrow_schema::extension::ExtensionType;
20use arrow_schema::extension::Uuid as ArrowUuid;
21use vortex_array::ArrayRef;
22use vortex_array::ExecutionCtx;
23use vortex_array::IntoArray;
24use vortex_array::arrays::ExtensionArray;
25use vortex_array::arrays::FixedSizeListArray;
26use vortex_array::arrays::PrimitiveArray;
27use vortex_array::arrays::extension::ExtensionArrayExt;
28use vortex_array::buffer::BufferHandle;
29use vortex_array::dtype::DType;
30use vortex_array::dtype::Nullability;
31use vortex_array::dtype::PType;
32use vortex_array::dtype::extension::ExtDType;
33use vortex_array::dtype::extension::ExtVTable;
34use vortex_array::extension::uuid::Uuid;
35use vortex_array::extension::uuid::UuidMetadata;
36use vortex_array::validity::Validity;
37use vortex_buffer::Alignment;
38use vortex_buffer::Buffer;
39use vortex_error::VortexExpect;
40use vortex_error::VortexResult;
41use vortex_session::registry::CachedId;
42use vortex_session::registry::Id;
43
44use crate::ArrowExport;
45use crate::ArrowExportVTable;
46use crate::ArrowImport;
47use crate::ArrowImportVTable;
48use crate::ArrowSession;
49use crate::ArrowSessionExt;
50use crate::nulls;
51use crate::session::has_valid_extension_type;
52
53const UUID_BYTE_LEN: i32 = 16;
54
55static ARROW_UUID: CachedId = CachedId::new(ArrowUuid::NAME);
56
57impl ArrowExportVTable for Uuid {
58    fn arrow_ext_id(&self) -> Id {
59        *ARROW_UUID
60    }
61
62    fn vortex_id(&self) -> Id {
63        Uuid.id()
64    }
65
66    // Encode all of these.
67    fn to_arrow_field(
68        &self,
69        name: &str,
70        dtype: &DType,
71        _session: &ArrowSession,
72    ) -> VortexResult<Option<Field>> {
73        let mut field = Field::new(
74            name.to_string(),
75            DataType::FixedSizeBinary(UUID_BYTE_LEN),
76            dtype.is_nullable(),
77        );
78        field
79            .try_with_extension_type(ArrowUuid)
80            .vortex_expect("FixedSizeBinary[16] is correct type for ArrowUuid");
81        Ok(Some(field))
82    }
83
84    fn execute_arrow(
85        &self,
86        array: ArrayRef,
87        _target: &Field,
88        ctx: &mut ExecutionCtx,
89    ) -> VortexResult<ArrowExport> {
90        let is_uuid = array
91            .dtype()
92            .as_extension_opt()
93            .map(|ext| ext.is::<Uuid>())
94            .unwrap_or(false);
95        if !is_uuid {
96            return Ok(ArrowExport::Unsupported(array));
97        }
98        Ok(ArrowExport::Exported(try_fsl_to_fsb(array, ctx)?))
99    }
100}
101
102impl ArrowImportVTable for Uuid {
103    fn arrow_ext_id(&self) -> Id {
104        *ARROW_UUID
105    }
106
107    fn from_arrow_field(
108        &self,
109        field: &Field,
110        _session: &ArrowSession,
111    ) -> VortexResult<Option<DType>> {
112        if !has_valid_extension_type::<ArrowUuid>(field) {
113            return Ok(None);
114        }
115
116        let storage_dtype = DType::FixedSizeList(
117            Arc::new(DType::Primitive(PType::U8, Nullability::NonNullable)),
118            UUID_BYTE_LEN as u32,
119            field.is_nullable().into(),
120        );
121
122        Ok(Some(DType::Extension(
123            ExtDType::try_with_vtable(Uuid, UuidMetadata::default(), storage_dtype)?.erased(),
124        )))
125    }
126
127    fn from_arrow_array(
128        &self,
129        array: ArrowArrayRef,
130        _field: &Field,
131        dtype: &DType,
132        _session: &ArrowSession,
133    ) -> VortexResult<ArrowImport> {
134        let DType::Extension(dtype) = dtype else {
135            return Ok(ArrowImport::Unsupported(array));
136        };
137        if !matches!(array.data_type(), DataType::FixedSizeBinary(UUID_BYTE_LEN))
138            || !dtype.is::<Uuid>()
139        {
140            return Ok(ArrowImport::Unsupported(array));
141        }
142
143        let fsb = array.as_fixed_size_binary();
144        let buffer = Buffer::from_arrow_buffer(fsb.values().clone(), Alignment::none());
145        let u8_array = PrimitiveArray::from_buffer_handle(
146            BufferHandle::new_host(buffer),
147            PType::U8,
148            Validity::NonNullable,
149        );
150        let validity = nulls(fsb.nulls(), dtype.is_nullable())?;
151
152        let storage = FixedSizeListArray::new(
153            u8_array.into_array(),
154            fsb.value_length() as u32,
155            validity,
156            fsb.len(),
157        )
158        .into_array();
159        Ok(ArrowImport::Imported(
160            ExtensionArray::new(dtype.clone(), storage).into_array(),
161        ))
162    }
163}
164
165fn try_fsl_to_fsb(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<ArrowArrayRef> {
166    let executed = array.execute::<ExtensionArray>(ctx)?;
167    let storage = executed.storage_array().clone();
168    let storage_arrow_type = DataType::FixedSizeList(
169        Arc::new(Field::new("item", DataType::UInt8, false)),
170        UUID_BYTE_LEN,
171    );
172
173    let storage_field = Field::new(
174        String::new(),
175        storage_arrow_type,
176        storage.dtype().is_nullable(),
177    );
178
179    let session = ctx.session().clone();
180    let arrow_storage = session
181        .arrow()
182        .execute_arrow(storage, Some(&storage_field), ctx)?;
183
184    let fsl = arrow_storage.as_fixed_size_list();
185    let bytes = fsl
186        .values()
187        .as_primitive::<UInt8Type>()
188        .values()
189        .inner()
190        .clone();
191
192    Ok(Arc::new(FixedSizeBinaryArray::new(
193        fsl.value_length(),
194        bytes,
195        fsl.nulls().cloned(),
196    )))
197}
vortex_arrow/uuid.rs

vortex_arrow/
uuid.rs
