Skip to content

Commit c546220

Browse files
myrrcmyrrc
authored
duckdb: UUID export support (#9494)
Support export of UUID in duckdb and UUID as a scalar in filter pushdown. Resolves: #9488 Signed-off-by: Mikhail Kot <mikhail@spiraldb.com>
1 parent d5ce13b commit c546220

9 files changed

Lines changed: 364 additions & 2 deletions

File tree

‎vortex-duckdb/src/convert/dtype.rs‎

Lines changed: 47 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,8 @@ use vortex::extension::datetime::TemporalMetadata;
5757
use vortex::extension::datetime::Time;
5858
use vortex::extension::datetime::TimeUnit;
5959
use vortex::extension::datetime::Timestamp;
60+
use vortex::extension::uuid::Uuid;
61+
use vortex::extension::uuid::UuidMetadata;
6062
use vortex_spatial::extension::LineString;
6163
use vortex_spatial::extension::MultiLineString;
6264
use vortex_spatial::extension::MultiPoint;
@@ -180,11 +182,18 @@ impl FromLogicalType for DType {
180182
)
181183
}
182184
DUCKDB_TYPE::DUCKDB_TYPE_VARIANT => DType::Variant(nullability),
185+
DUCKDB_TYPE::DUCKDB_TYPE_UUID => DType::Extension(
186+
ExtDType::<Uuid>::try_with_vtable(
187+
Uuid,
188+
UuidMetadata::default(),
189+
uuid_storage_dtype(nullability),
190+
)?
191+
.erased(),
192+
),
183193
other @ (DUCKDB_TYPE::DUCKDB_TYPE_TIME_TZ
184194
| DUCKDB_TYPE::DUCKDB_TYPE_INTERVAL
185195
| DUCKDB_TYPE::DUCKDB_TYPE_ENUM
186196
| DUCKDB_TYPE::DUCKDB_TYPE_MAP
187-
| DUCKDB_TYPE::DUCKDB_TYPE_UUID
188197
| DUCKDB_TYPE::DUCKDB_TYPE_UNION
189198
| DUCKDB_TYPE::DUCKDB_TYPE_BIT
190199
| DUCKDB_TYPE::DUCKDB_TYPE_ANY
@@ -254,6 +263,10 @@ impl TryFrom<&DType> for LogicalType {
254263
return temporal_to_duckdb(temporal);
255264
}
256265

266+
if ext_dtype.is::<Uuid>() {
267+
return Ok(LogicalType::new(DUCKDB_TYPE::DUCKDB_TYPE_UUID));
268+
}
269+
257270
// Native geometry types and WKB all surface to DuckDB as GEOMETRY so `ST_*` bind.
258271
if let Some(spatial_metadata) = ext_dtype
259272
.metadata_opt::<Point>()
@@ -275,6 +288,14 @@ impl TryFrom<&DType> for LogicalType {
275288
}
276289
}
277290

291+
fn uuid_storage_dtype(nullability: Nullability) -> DType {
292+
DType::FixedSizeList(
293+
Arc::new(DType::Primitive(U8, Nullability::NonNullable)),
294+
16,
295+
nullability,
296+
)
297+
}
298+
278299
fn temporal_to_duckdb(temporal: TemporalMetadata) -> VortexResult<LogicalType> {
279300
let duckdb_type = match temporal {
280301
TemporalMetadata::Timestamp(unit, None) => match unit {
@@ -637,6 +658,31 @@ mod tests {
637658
Ok(())
638659
}
639660

661+
#[rstest]
662+
#[case(Nullability::NonNullable)]
663+
#[case(Nullability::Nullable)]
664+
fn test_uuid_roundtrip(#[case] nullability: Nullability) -> VortexResult<()> {
665+
use vortex::extension::uuid::Uuid;
666+
use vortex::extension::uuid::UuidMetadata;
667+
668+
let storage = DType::FixedSizeList(
669+
Arc::new(DType::Primitive(PType::U8, Nullability::NonNullable)),
670+
16,
671+
nullability,
672+
);
673+
let vortex_uuid = DType::Extension(
674+
ExtDType::<Uuid>::try_with_vtable(Uuid, UuidMetadata::default(), storage)?.erased(),
675+
);
676+
677+
let duckdb_uuid = LogicalType::try_from(&vortex_uuid)?;
678+
assert_eq!(duckdb_uuid.as_type_id(), cpp::DUCKDB_TYPE::DUCKDB_TYPE_UUID);
679+
680+
let original = DType::from_logical_type(&duckdb_uuid, nullability)?;
681+
assert_eq!(original, vortex_uuid);
682+
683+
Ok(())
684+
}
685+
640686
#[test]
641687
fn test_unsupported_extension_type() {
642688
#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)]

‎vortex-duckdb/src/convert/scalar.rs‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
use vortex::array::match_each_native_simd_ptype;
2323
use vortex::dtype::DType;
2424
use vortex::dtype::DecimalDType;
25+
use vortex::dtype::Nullability::NonNullable;
2526
use vortex::dtype::Nullability::Nullable;
2627
use vortex::dtype::PType;
2728
use vortex::dtype::PType::I32;
@@ -30,6 +31,7 @@ use vortex::dtype::half::f16;
3031
use vortex::error::VortexError;
3132
use vortex::error::VortexResult;
3233
use vortex::error::vortex_bail;
34+
use vortex::error::vortex_ensure;
3335
use vortex::error::vortex_err;
3436
use vortex::extension::datetime::AnyTemporal;
3537
use vortex::extension::datetime::Date;
@@ -38,6 +40,7 @@ use vortex::extension::datetime::Time;
3840
use vortex::extension::datetime::TimeUnit;
3941
use vortex::extension::datetime::Timestamp;
4042
use vortex::extension::datetime::TimestampOptions;
43+
use vortex::extension::uuid::Uuid;
4144
use vortex::scalar::BinaryScalar;
4245
use vortex::scalar::BoolScalar;
4346
use vortex::scalar::DecimalScalar;
@@ -315,6 +318,19 @@ impl<'a> TryFrom<&'a ValueRef> for Scalar {
315318
ext.clone(),
316319
Scalar::binary(b, Nullable),
317320
)),
321+
DType::Extension(ext) if ext.is::<Uuid>() => {
322+
vortex_ensure!(b.len() == 16, "UUID blob must be 16 bytes, got {}", b.len());
323+
let children = b
324+
.iter()
325+
.map(|&byte| Scalar::primitive(byte, NonNullable))
326+
.collect();
327+
let storage = Scalar::fixed_size_list(
328+
DType::Primitive(PType::U8, NonNullable),
329+
children,
330+
Nullable,
331+
);
332+
Ok(Scalar::extension_ref(ext.clone(), storage))
333+
}
318334
_ => vortex_bail!("Cannot convert DuckDB blob to Vortex scalar of dtype {dtype}"),
319335
},
320336
ExtractedValue::Date(days) => Ok(Scalar::extension::<Date>(

‎vortex-duckdb/src/convert/vector.rs‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@ use vortex::error::VortexExpect;
2929
use vortex::error::VortexResult;
3030
use vortex::error::vortex_bail;
3131
use vortex::extension::datetime::TimeUnit;
32+
use vortex::extension::uuid::Uuid;
33+
use vortex::extension::uuid::UuidMetadata;
3234
use vortex::mask::Mask;
3335
use vortex_spatial::extension::SpatialMetadata;
3436
use vortex_spatial::extension::WellKnownBinary;
@@ -316,6 +318,29 @@ pub fn flat_vector_to_vortex(vector: &VectorRef, len: usize) -> VortexResult<Arr
316318
}
317319
.map(|a| a.into_array())
318320
}
321+
DUCKDB_TYPE::DUCKDB_TYPE_UUID => {
322+
let data = vector.as_slice_with_len::<i128>(len);
323+
let mut bytes = BufferMut::<u8>::with_capacity(len * 16);
324+
for v in data {
325+
let be = (*v as u128) ^ (1u128 << 127);
326+
bytes.extend_from_slice(&be.to_be_bytes());
327+
}
328+
329+
let storage = FixedSizeListArray::try_new(
330+
PrimitiveArray::new(bytes.freeze(), Validity::NonNullable).into_array(),
331+
16,
332+
vector.validity_ref(len).to_validity(),
333+
len,
334+
)?;
335+
let ext_dtype = ExtDType::<Uuid>::try_with_vtable(
336+
Uuid,
337+
UuidMetadata::default(),
338+
storage.dtype().clone(),
339+
)?
340+
.erased();
341+
342+
Ok(ExtensionArray::try_new(ext_dtype, storage.into_array())?.into_array())
343+
}
319344
DUCKDB_TYPE::DUCKDB_TYPE_ARRAY => {
320345
let array_elem_size = vector.logical_type().array_type_array_size();
321346
let child_data = flat_vector_to_vortex(

‎vortex-duckdb/src/duckdb/value.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,7 @@ impl ValueRef {
9999
unsafe { cpp::duckdb_free(ptr.cast()) };
100100
ExtractedValue::Varchar(string)
101101
}
102-
DUCKDB_TYPE::DUCKDB_TYPE_BLOB => {
102+
DUCKDB_TYPE::DUCKDB_TYPE_BLOB | DUCKDB_TYPE::DUCKDB_TYPE_UUID => {
103103
ExtractedValue::Blob(unsafe { take_blob(cpp::duckdb_get_blob(self.as_ptr())) })
104104
}
105105
DUCKDB_TYPE::DUCKDB_TYPE_GEOMETRY => {

‎vortex-duckdb/src/exporter/extension.rs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ use vortex::array::arrays::extension::ExtensionArrayExt;
88
use vortex::array::extension::datetime::AnyTemporal;
99
use vortex::error::VortexResult;
1010
use vortex::error::vortex_bail;
11+
use vortex::extension::uuid::Uuid;
1112
use vortex_spatial::extension::LineString;
1213
use vortex_spatial::extension::LineStringData;
1314
use vortex_spatial::extension::MultiLineString;
@@ -26,6 +27,7 @@ use vortex_spatial::extension::WellKnownBinaryData;
2627
use crate::exporter::ColumnExporter;
2728
use crate::exporter::spatial;
2829
use crate::exporter::temporal;
30+
use crate::exporter::uuid;
2931

3032
pub(crate) fn new_exporter(
3133
ext: ExtensionArray,
@@ -35,6 +37,10 @@ pub(crate) fn new_exporter(
3537
return temporal::new_exporter(TemporalArray::try_from(ext)?, ctx);
3638
}
3739

40+
if ext.ext_dtype().is::<Uuid>() {
41+
return uuid::new_exporter(ext, ctx);
42+
}
43+
3844
if ext.ext_dtype().is::<WellKnownBinary>() {
3945
return spatial::new_wkb_exporter(WellKnownBinaryData::try_from(ext)?, ctx);
4046
}

‎vortex-duckdb/src/exporter/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ mod sequence;
2020
mod spatial;
2121
mod struct_;
2222
mod temporal;
23+
mod uuid;
2324
mod validity;
2425
mod varbinview;
2526
mod vector;

0 commit comments

Comments
 (0)