Skip to content

Commit

Permalink
feat: create union type for virtual tables (#269)
Browse files Browse the repository at this point in the history
* add ability to create union type

* rework to use proper union functions

* check for panic

* swap for if

* tweak import

* remove workaround

* remove deref

* Remove duplicate use
  • Loading branch information
Mause authored Mar 21, 2024
1 parent 62089b6 commit 34a6448
Showing 1 changed file with 50 additions and 4 deletions.
54 changes: 50 additions & 4 deletions src/vtab/logical_type.rs
Original file line number Diff line number Diff line change
Expand Up @@ -217,6 +217,23 @@ impl LogicalType {
}
}

/// Make a `LogicalType` for `union`
pub fn union_type(fields: &[(&str, LogicalType)]) -> Self {
let keys: Vec<CString> = fields.iter().map(|f| CString::new(f.0).unwrap()).collect();
let values: Vec<duckdb_logical_type> = fields.iter().map(|it| it.1.ptr).collect();
let name_ptrs = keys.iter().map(|it| it.as_ptr()).collect::<Vec<*const c_char>>();

unsafe {
Self {
ptr: duckdb_create_union_type(
values.as_slice().as_ptr().cast_mut(),
name_ptrs.as_slice().as_ptr().cast_mut(),
fields.len() as idx_t,
),
}
}
}

/// Logical type ID
pub fn id(&self) -> LogicalTypeId {
let duckdb_type_id = unsafe { duckdb_get_type_id(self.ptr) };
Expand All @@ -227,16 +244,22 @@ impl LogicalType {
pub fn num_children(&self) -> usize {
match self.id() {
LogicalTypeId::Struct => unsafe { duckdb_struct_type_child_count(self.ptr) as usize },
LogicalTypeId::Union => unsafe { duckdb_union_type_member_count(self.ptr) as usize },
LogicalTypeId::List => 1,
_ => 0,
}
}

/// Logical type child name by idx
///
/// Panics if the logical type is not a struct or union
pub fn child_name(&self, idx: usize) -> String {
assert_eq!(self.id(), LogicalTypeId::Struct);
unsafe {
let child_name_ptr = duckdb_struct_type_child_name(self.ptr, idx as u64);
let child_name_ptr = match self.id() {
LogicalTypeId::Struct => duckdb_struct_type_child_name(self.ptr, idx as u64),
LogicalTypeId::Union => duckdb_union_type_member_name(self.ptr, idx as u64),
_ => panic!("not a struct or union"),
};
let c_str = CString::from_raw(child_name_ptr);
let name = c_str.to_str().unwrap();
name.to_string()
Expand All @@ -245,14 +268,20 @@ impl LogicalType {

/// Logical type child by idx
pub fn child(&self, idx: usize) -> Self {
let c_logical_type = unsafe { duckdb_struct_type_child_type(self.ptr, idx as u64) };
let c_logical_type = unsafe {
match self.id() {
LogicalTypeId::Struct => duckdb_struct_type_child_type(self.ptr, idx as u64),
LogicalTypeId::Union => duckdb_union_type_member_type(self.ptr, idx as u64),
_ => panic!("not a struct or union"),
}
};
Self::from(c_logical_type)
}
}

#[cfg(test)]
mod test {
use crate::vtab::LogicalType;
use super::{LogicalType, LogicalTypeId};

#[test]
fn test_struct() {
Expand Down Expand Up @@ -280,4 +309,21 @@ mod test {
assert_eq!(typ.decimal_width(), 0);
assert_eq!(typ.decimal_scale(), 0);
}

#[test]
fn test_union_type() {
let fields = &[
("hello", LogicalType::new(LogicalTypeId::Boolean)),
("world", LogicalType::new(LogicalTypeId::Integer)),
];
let typ = LogicalType::union_type(fields);

assert_eq!(typ.num_children(), 2);

assert_eq!(typ.child_name(0), "hello");
assert_eq!(typ.child(0).id(), LogicalTypeId::Boolean);

assert_eq!(typ.child_name(1), "world");
assert_eq!(typ.child(1).id(), LogicalTypeId::Integer);
}
}

0 comments on commit 34a6448

Please sign in to comment.